ovpn-lb-socks5/openvpn.py
2020-04-03 15:50:29 -04:00

139 lines
4.1 KiB
Python

"""
openvpn supervisor
"""
import asyncio
import tempfile
import os
import time
import openvpn_api
def generate_config(in_fp, cfg):
def change_item(config, item_key, item_content):
output_lines = []
changed = False
for i in config.splitlines():
if i.startswith(item_key+" ") or i == item_key:
output_lines.append(f"{item_key} {item_content}")
changed = True
else:
output_lines.append(i)
if not changed:
output_lines.append(f"{item_key} {item_content}")
return "\n".join(output_lines)
# _, path = tempfile.mkstemp()
with open(in_fp, "r") as in_f:
config = in_f.read()
for key, value in cfg.items():
config = change_item(config, key, value)
return config
IDLE = "idle"
RUNNING = "running"
class Openvpn:
def __init__(self, cfg_fp, interface, folder_path, management_port, name=None, additional_cfg={}, loop=None):
self.cfg_fp = cfg_fp
self.interface = interface
self.folder_path = folder_path
self.management_port = management_port
self.PID = 0
self.status = IDLE
self.proc = None
self.exit_future = None
self.additional_cfg = additional_cfg
self.run_task = None
self.openvpn_api = None
self.name = name
if loop:
self.loop = loop
else:
self.loop = asyncio.get_event_loop()
def get_cfg(self):
self.log_fp = os.path.join(self.folder_path, "log.txt")
cfg = {
"dev": self.interface,
"dev-type": "tun", # TODO: add code to read dev-type
"management": f"localhost {self.management_port}",
"log-append": self.log_fp
}
cfg.update(self.additional_cfg)
return cfg
def generate_config_file(self):
cfg = self.get_cfg()
self.config_fp = os.path.join(self.folder_path, "cfg.txt")
config = generate_config(self.cfg_fp, cfg)
with open(self.config_fp, "w") as config_f:
config_f.write(config)
return self.config_fp
def start(self):
if self.status == IDLE:
self.status = RUNNING
config_fp = self.generate_config_file()
cmd = " ".join(["openvpn", "--config", config_fp])
self.run_task = self.loop.create_task(self.run(cmd))
def get_log(self):
# regenerate log_fp
_ = self.get_cfg()
try:
with open(self.log_fp, "r") as log_f:
return log_f.read()
except:
return ""
def clear_log(self):
os.remove(self.log_fp)
async def stop(self):
if self.status == RUNNING:
try:
await self.proc.kill()
except:
print("kill failed")
self.status = IDLE
self.run_task.cancel()
async def restart(self):
await self.stop()
self.start()
async def run(self, cmd):
print(f"run: {cmd}")
self.exit_future = asyncio.Future()
while self.status == RUNNING:
print("create proc")
print(self.status)
proc = await asyncio.create_subprocess_shell(
cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE)
self.proc = proc
self.PID = proc.pid
self.openvpn_api = openvpn_api.VPN(
'localhost', int(self.management_port))
print(f"pid: {proc.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)
self.exit_future.set_result(True)
self.status = IDLE
if __name__ == "__main__":
folder_fp = "/home/mantao/Desktop/t/"
cfg_fp = "/home/mantao/Desktop/t/TCP_Files/UK2-TCP.ovpn"
o1 = Openvpn(cfg_fp, "tun0", folder_fp, 8001, additional_cfg={
"auth-user-pass": "/home/mantao/Desktop/t/fast.txt"})
o1.start()