connect2socks_proxy.py

#!/usr/bin/env python3

from argparse import ArgumentParser
from functools import partial
from shutil import copy
from subprocess import Popen, PIPE
from shlex import split
import os
import re
import atexit
import sys
from time import sleep

SOCKS5_PORT = 1337
REVERSE_SSH_PORT = 2222

def run_cmd(cmd: str, wait=True, stdout=PIPE, stderr=PIPE) -> Popen:
    print(cmd)
    args = split(cmd)
    proc = Popen(args=args, stdout=stdout, stderr=stderr)

    if wait:
        proc.wait()

    return proc

def is_port_in_use(port: int) -> bool:
    import socket
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('localhost', port)) == 0

def setup():
    if not os.path.exists("/usr/local/bin/tun2socks"):
        curdir = os.curdir

        run_cmd("apt update")
        run_cmd("apt install make git wget")

        # install golang
        run_cmd("wget -q https://go.dev/dl/go1.19.2.linux-amd64.tar.gz -O /tmp/golang.tar.gz")
        run_cmd("rm -rf /usr/local/go")
        run_cmd("tar -C /usr/local -xzf /tmp/golang.tar.gz")
        os.environ["PATH"] = os.environ["PATH"] + os.pathsep + "/usr/local/go/bin"

        # Install tun2socks
        run_cmd("git clone https://github.com/xjasonlyu/tun2socks.git --depth 1")
        os.chdir("tun2socks")
        run_cmd("make tun2socks -j")

        copy("./build/tun2socks", "/usr/local/bin")

        os.chdir(curdir)

def get_current_interface():
    routes = run_cmd("ip route").communicate()[0].decode("utf8").splitlines()
    iface = None

    for route in routes:
        if "10.10.10.10" in route:
            continue

        if "default" in route:
            iface = re.search(string=route, pattern=r".*dev (e[0-9,a-z]*) .*").group(1)
            break

    if iface is None:
        raise RuntimeError(f"Cannot find inteface! {routes}")

    return iface
            

def connect(username: str, ip: str, port: int):
    socks5_proxy_process = run_cmd(f"ssh -D {SOCKS5_PORT} {username}@{ip} -N -p {port}", wait=False)
    for i in range(10):
        print(f"Waiting for opening port {SOCKS5_PORT}. Attemt #{i}")
        sleep(1)
        if is_port_in_use(SOCKS5_PORT):
            print(f"Success, connection is established")
            break
    else:
        print(f"Can't connect! Port {SOCKS5_PORT} is closed! Exiting.")
        sys.exit(-1)

    run_cmd("ip tuntap add mode tun dev tun0")
    run_cmd("ip addr add 10.10.10.10/24 dev tun0")
    run_cmd("ip link set dev tun0 up")

    interface = get_current_interface()
    cmd = f"/usr/local/bin/tun2socks -device tun0 -proxy socks5://localhost:{SOCKS5_PORT} -interface {interface}"
    tun2socks_process = run_cmd(cmd, wait=False)
    run_cmd("ip route add default via 10.10.10.10")

    return [socks5_proxy_process, tun2socks_process]

def disconnect(processes):
    for process in processes:
        process: Popen
        process.kill()
    
    routes = run_cmd("ip route").communicate()[0].decode("utf8")
    if "10.10.10.10" in routes:
        run_cmd("ip route delete default via 10.10.10.10")
    
    links = run_cmd("ip link").communicate()[0].decode("utf8")
    if "tun0" in links:
        run_cmd("ip link delete tun0")

def main():
    parser = ArgumentParser()
    parser.add_argument("--user", required=True, type=str)
    parser.add_argument("--port", type=int, default=22)
    parser.add_argument("--ip", type=str, default="localhost")

    args = parser.parse_args()

    setup()
    disconnect([])
    processes = connect(args.user, args.ip, args.port)

    atexit.register(partial(disconnect, processes))

    try:
        for p in processes:
            p.wait()
    except:
        disconnect(processes)
        return -1

main()