diff --git a/o_manager.py b/o_manager.py index 79cdfab..73b6847 100644 --- a/o_manager.py +++ b/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__": diff --git a/test.py b/test.py index cbfbaa0..0be4cff 100644 --- a/test.py +++ b/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():