#!/usr/bin/env python3
"""
SOCKS5 转 HTTP 代理转发工具
将 SSH 隧道的 SOCKS5 代理 (1080) 转为 HTTP 代理 (8080)
供 CCXT 等只支持 HTTP 代理的库使用
"""

import socket
import threading
import select
import sys
from urllib.parse import urlparse


class Socks2Http:
    def __init__(self, socks_host='127.0.0.1', socks_port=1080, http_port=8080):
        self.socks_host = socks_host
        self.socks_port = socks_port
        self.http_port = http_port
        self.server = None
        self.running = False
        
    def start(self):
        """启动 HTTP 代理服务器"""
        self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server.bind(('127.0.0.1', self.http_port))
        self.server.listen(5)
        self.running = True
        
        print(f"HTTP 代理已启动: http://127.0.0.1:{self.http_port}")
        print(f"转发到 SOCKS5: {self.socks_host}:{self.socks_port}")
        
        while self.running:
            try:
                client, addr = self.server.accept()
                thread = threading.Thread(target=self._handle_client, args=(client,))
                thread.daemon = True
                thread.start()
            except Exception as e:
                if self.running:
                    print(f"接受连接错误: {e}")
                    
    def stop(self):
        """停止代理服务器"""
        self.running = False
        if self.server:
            self.server.close()
        print("HTTP 代理已停止")
        
    def _handle_client(self, client):
        """处理客户端连接"""
        try:
            # 读取 HTTP 请求
            request = client.recv(4096)
            if not request:
                client.close()
                return
                
            # 解析 CONNECT 请求
            first_line = request.split(b'\r\n')[0].decode('utf-8', errors='ignore')
            
            if first_line.startswith('CONNECT'):
                # HTTPS 代理
                parts = first_line.split()
                if len(parts) >= 2:
                    target = parts[1]
                    host, port = target.rsplit(':', 1)
                    port = int(port)
                    
                    # 连接到 SOCKS5 代理
                    remote = self._connect_via_socks5(host, port)
                    if remote:
                        client.send(b'HTTP/1.1 200 Connection established\r\n\r\n')
                        self._relay(client, remote)
                    else:
                        client.send(b'HTTP/1.1 502 Bad Gateway\r\n\r\n')
            else:
                # HTTP 代理
                lines = first_line.split()
                if len(lines) >= 2:
                    url = lines[1]
                    parsed = urlparse(url)
                    host = parsed.hostname
                    port = parsed.port or 80
                    
                    # 连接到 SOCKS5 代理
                    remote = self._connect_via_socks5(host, port)
                    if remote:
                        # 修改请求，去掉完整 URL
                        new_request = request.replace(
                            f'{lines[0]} {url} '.encode(),
                            f'{lines[0]} {parsed.path or "/"}{"?" + parsed.query if parsed.query else ""} '.encode()
                        )
                        remote.send(new_request)
                        self._relay(client, remote)
                        
        except Exception as e:
            print(f"处理客户端错误: {e}")
        finally:
            client.close()
            
    def _connect_via_socks5(self, host, port):
        """通过 SOCKS5 代理连接目标服务器"""
        try:
            # 连接到 SOCKS5 代理
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(30)
            sock.connect((self.socks_host, self.socks_port))
            
            # SOCKS5 握手
            # 1. 发送认证方法
            sock.send(b'\x05\x01\x00')  # VER=5, NMETHODS=1, METHOD=0 (无认证)
            resp = sock.recv(2)
            if resp[0] != 0x05 or resp[1] != 0x00:
                sock.close()
                return None
                
            # 2. 发送连接请求
            req = b'\x05\x01\x00\x03'  # VER=5, CMD=CONNECT, RSV=0, ATYP=DOMAIN
            req += bytes([len(host)]) + host.encode()
            req += bytes([(port >> 8) & 0xFF, port & 0xFF])
            sock.send(req)
            
            # 3. 读取响应
            resp = sock.recv(10)
            if len(resp) < 4 or resp[1] != 0x00:
                sock.close()
                return None
                
            return sock
            
        except Exception as e:
            print(f"SOCKS5 连接错误: {e}")
            return None
            
    def _relay(self, client, remote):
        """双向转发数据"""
        try:
            while True:
                readable, _, _ = select.select([client, remote], [], [], 60)
                if not readable:
                    break
                    
                if client in readable:
                    data = client.recv(4096)
                    if not data:
                        break
                    remote.send(data)
                    
                if remote in readable:
                    data = remote.recv(4096)
                    if not data:
                        break
                    client.send(data)
        except:
            pass
        finally:
            client.close()
            remote.close()


if __name__ == '__main__':
    proxy = Socks2Http(socks_port=1080, http_port=8080)
    try:
        proxy.start()
    except KeyboardInterrupt:
        proxy.stop()
