add load balancer management

o_manager.py

@@ -1,22 +1,30 @@
 import os
 from openvpn import Openvpn
 import asyncio
+from tornado.template import Template
+import signal
 
 
 class OManager:
-    def __init__(self, base_folder, base_port=8001):
+    def __init__(self, base_folder, base_port=8001, loop=None):
         self.base_folder = base_folder
         self.base_port = base_port
         self.ops = []
         self.idx = 0
+        self.running_idx = []
+        self.run_task = []
+        if loop:
+            self.loop = loop
+        else:
+            self.loop = asyncio.get_event_loop()
 
     def new_op(self, cfg_fp, name=None, additional_cfg={}):
-        interface = f"tun{self.idx}"
         folder_path = os.path.join(self.base_folder, f"session{self.idx}")
         if not os.path.isdir(folder_path):
             os.makedirs(folder_path)
         if not name:
             name = f"openvpn-{self.idx}"
+        os.system(f"groupadd vpn{self.idx}")
         op = Openvpn(cfg_fp, self.idx, folder_path,
                      f"{self.base_port + self.idx}", "script.sh.template", "3proxy.cfg.template", name=name,
                      additional_cfg=additional_cfg)
@@ -24,16 +32,82 @@ class OManager:
         self.idx += 1
         return op
 
+    def generate_lb_cfg(self):
+        lb_cfg_fp = os.path.join(
+            self.base_folder, "go-socks-lb.yml")
+        with open("go-socks-lb.yml.template", "r") as template_f:
+            buf = template_f.read()
+        template = Template(buf)
+        items = {
+            "running_idx": self.running_idx,
+            "port": f"{1080 + self.idx}"
+        }
+        lb_cfg = template.generate(**items)
+        with open(lb_cfg_fp, "wb") as lb_cfg_f:
+            lb_cfg_f.write(lb_cfg)
+        return lb_cfg_fp
+
     def start_op(self, idx):
         self.ops[idx].start()
+        if idx not in self.running_idx:
+            self.running_idx.append(idx)
+        self.reset_lb()
+
+    def stop_op(self, idx):
+        self.loop.create_task(self.ops[idx].stop())
+        if idx in self.running_idx:
+            self.running_idx.remove(idx)
+        self.reset_lb()
 
     def get_all_ops(self):
         return self.ops
 
+    def start_all(self):
+        for op in self.ops:
+            op.start()
+            idx = op.idx
+            if idx not in self.running_idx:
+                self.running_idx.append(idx)
+
     def stop_all(self):
-        loop = asyncio.get_event_loop()
+        loop = self.loop
         for op in self.ops:
             loop.create_task(op.stop())
+        self.running_idx = []
+        self.reset_lb()
+
+    def reset_lb(self):
+        try:
+            os.kill(self.PID, signal.SIGINT)
+        except Exception as err:
+            print("kill failed:", err)
+        for task in self.run_task:
+            task.cancel()
+        self.run_task = []
+        lb_cfg_fp = self.generate_lb_cfg()
+        # TODO: remove --mute-replay-warnings
+        lb_cmd = " ".join(["go-socks-lb/go-socks-lb", "-config",
+                           lb_cfg_fp, "-bind", "0.0.0.0:7000"])
+        self.run_task.append(self.loop.create_task(
+            self.run_cmd(lb_cmd)))
+
+    async def run_cmd(self, cmd):
+        while True:
+            proc = await asyncio.create_subprocess_exec(
+                cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE)
+            print("started")
+            self.proc = proc
+            self.PID = proc.pid
+            print(f"pid: {self.PID}")
+            stdout, stderr = await proc.communicate()
+            print(f'[{cmd!r} exited with {proc.returncode}]')
+            if stdout:
+                print(f'[stdout]\n{stdout.decode()}')
+            if stderr:
+                print(f'[stderr]\n{stderr.decode()}')
+            await asyncio.sleep(5)
 
 
 if __name__ == "__main__":

test.py

@@ -66,12 +66,19 @@ class StartInstatnceHandler(tornado.web.RequestHandler):
             om.start_op(int(idx))
             self.write("start sucess!")
         else:
-            self.write("need idx")
+            om.start_all()
+            self.write("start all")
 
 
 class StopAll(tornado.web.RequestHandler):
     def get(self):
-        om.stop_all()
+        idx = self.get_query_argument("i", None)
+        if idx != None:
+            om.stop_op(int(idx))
+            self.write("stop sucess!")
+        else:
+            om.stop_all()
+            self.write("stop all")
 
 
 def make_app():