dots.ocr/zhch/secure_tunnel.py

#!/usr/bin/env python3
"""
安全的端口转发 - 只允许特定进程访问
"""
import subprocess
import psutil
import time
import os
import signal
import sys
import socket
import threading
import json
from pathlib import Path
from http.server import HTTPServer, BaseHTTPRequestHandler
import requests

class SecureProxyHandler(BaseHTTPRequestHandler):
    """安全代理处理器 - 只允许特定进程访问"""
    
    def __init__(self, allowed_pids, target_host, target_port, *args, **kwargs):
        self.allowed_pids = allowed_pids
        self.target_host = target_host
        self.target_port = target_port
        super().__init__(*args, **kwargs)
    
    def is_request_allowed(self):
        """检查请求是否来自允许的进程"""
        try:
            # 获取客户端连接信息
            client_ip = self.client_address[0]
            
            # 如果不是本地连接，直接拒绝
            if client_ip not in ['127.0.0.1', '::1']:
                return False, f"Non-local connection from {client_ip}"
            
            # 检查当前运行的进程
            current_pids = set()
            for proc in psutil.process_iter(['pid', 'name', 'cmdline']):
                try:
                    proc_info = proc.info
                    proc_name = proc_info['name']
                    cmdline = ' '.join(proc_info['cmdline'] or [])
                    
                    # 检查是否是允许的进程
                    for allowed_process in ['Code', 'github-copilot', 'vllm', 'python']:
                        if (allowed_process.lower() in proc_name.lower() or 
                            allowed_process.lower() in cmdline.lower()):
                            current_pids.add(proc.info['pid'])
                            
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
            
            # 检查是否有允许的进程在运行
            if not current_pids:
                return False, "No allowed processes running"
            
            return True, "Request allowed"
            
        except Exception as e:
            return False, f"Error checking request: {e}"
    
    def proxy_request(self, method='GET', data=None):
        """代理请求到目标服务器"""
        allowed, reason = self.is_request_allowed()
        
        if not allowed:
            self.send_error(403, f"Access denied: {reason}")
            return
        
        try:
            # 构建目标 URL
            target_url = f"http://{self.target_host}:{self.target_port}{self.path}"
            
            # 准备请求头
            headers = {}
            for header, value in self.headers.items():
                if header.lower() not in ['host', 'connection']:
                    headers[header] = value
            
            # 发送请求
            if method == 'GET':
                response = requests.get(target_url, headers=headers, timeout=30)
            elif method == 'POST':
                response = requests.post(target_url, headers=headers, data=data, timeout=30)
            else:
                self.send_error(405, "Method not allowed")
                return
            
            # 返回响应
            self.send_response(response.status_code)
            for header, value in response.headers.items():
                if header.lower() not in ['connection', 'transfer-encoding']:
                    self.send_header(header, value)
            self.end_headers()
            
            self.wfile.write(response.content)
            
            # 记录访问日志
            self.log_message(f"Proxied {method} {self.path} -> {response.status_code}")
            
        except requests.RequestException as e:
            self.send_error(502, f"Proxy error: {e}")
        except Exception as e:
            self.send_error(500, f"Internal error: {e}")
    
    def do_GET(self):
        self.proxy_request('GET')
    
    def do_POST(self):
        content_length = int(self.headers.get('Content-Length', 0))
        post_data = self.rfile.read(content_length)
        self.proxy_request('POST', post_data)
    
    def log_message(self, format, *args):
        """自定义日志格式"""
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        sys.stderr.write(f"[{timestamp}] {format % args}\n")

class SecureTunnel:
    def __init__(self):
        self.config_file = Path.home() / '.secure_tunnel_config.json'
        self.proxy_server = None
        self.tunnel_process = None
        self.proxy_thread = None
        
    def load_config(self):
        """加载配置"""
        default_config = {
            "remote_host": "10.192.72.11",
            "remote_user": "ubuntu",
            # "ssh_key": str(Path.home() / ".ssh/id_dotsocr_tunnel"),
            "local_forward_port": 8082,
            "remote_service_port": 8101,
            "secure_proxy_port": 7281,  # 安全代理端口
            "target_proxy_host": "127.0.0.1",
            "target_proxy_port": 7890,
            "allowed_processes": ["Code", "github-copilot", "vllm"],
            "check_interval": 30
        }
        
        if self.config_file.exists():
            with open(self.config_file, 'r') as f:
                config = json.load(f)
                for key, value in default_config.items():
                    if key not in config:
                        config[key] = value
                return config
        else:
            with open(self.config_file, 'w') as f:
                json.dump(default_config, f, indent=2)
            return default_config
    
    def start_secure_proxy(self, config):
        """启动安全代理服务器"""
        def create_handler(*args, **kwargs):
            return SecureProxyHandler(
                config['allowed_processes'],
                config['target_proxy_host'],
                config['target_proxy_port'],
                *args, **kwargs
            )
        
        try:
            self.proxy_server = HTTPServer(('127.0.0.1', config['secure_proxy_port']), create_handler)
            print(f"🔒 Secure proxy started on port {config['secure_proxy_port']}")
            
            def run_server():
                self.proxy_server.serve_forever()
            
            self.proxy_thread = threading.Thread(target=run_server, daemon=True)
            self.proxy_thread.start()
            return True
            
        except Exception as e:
            print(f"❌ Failed to start secure proxy: {e}")
            return False
    
    def start_tunnel(self, config):
        """启动 SSH 隧道（使用安全代理端口）"""
        if self.tunnel_process and self.tunnel_process.poll() is None:
            return True
        
        ssh_cmd = [
            'ssh',
            # '-i', config['ssh_key'],
            '-o', 'ExitOnForwardFailure=yes',
            '-o', 'ServerAliveInterval=30',
            '-o', 'ServerAliveCountMax=3',
            '-L', f"{config['local_forward_port']}:localhost:{config['remote_service_port']}",
            '-R', f"7280:localhost:{config['secure_proxy_port']}",  # 转发到安全代理
            '-N',
            f"{config['remote_user']}@{config['remote_host']}"
        ]
        
        try:
            print(f"🚀 Starting secure SSH tunnel...")
            self.tunnel_process = subprocess.Popen(
                ssh_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                preexec_fn=os.setsid
            )
            
            time.sleep(3)
            if self.tunnel_process.poll() is None:
                print("✅ Secure SSH tunnel started")
                return True
            else:
                stdout, stderr = self.tunnel_process.communicate()
                print(f"❌ SSH tunnel failed: {stderr.decode()}")
                return False
                
        except Exception as e:
            print(f"❌ Error starting tunnel: {e}")
            return False
    
    def stop_all(self):
        """停止所有服务"""
        # 停止代理服务器
        if self.proxy_server:
            self.proxy_server.shutdown()
            self.proxy_server.server_close()
            print("🔒 Secure proxy stopped")
        
        # 停止 SSH 隧道
        if self.tunnel_process:
            try:
                os.killpg(os.getpgid(self.tunnel_process.pid), signal.SIGTERM)
                self.tunnel_process.wait(timeout=5)
                print("🚀 SSH tunnel stopped")
            except:
                try:
                    os.killpg(os.getpgid(self.tunnel_process.pid), signal.SIGKILL)
                except:
                    pass
    
    def run(self):
        """运行安全隧道"""
        config = self.load_config()
        
        print("🛡️ Starting Secure DotsOCR Tunnel")
        print("=" * 50)
        print(f"🔒 Secure proxy port: {config['secure_proxy_port']}")
        print(f"🎯 Target proxy: {config['target_proxy_host']}:{config['target_proxy_port']}")
        print(f"✅ Allowed processes: {', '.join(config['allowed_processes'])}")
        print("=" * 50)
        
        # 设置信号处理
        def signal_handler(signum, frame):
            print("\n🛑 Shutting down secure tunnel...")
            self.stop_all()
            sys.exit(0)
        
        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)
        
        try:
            # 启动安全代理
            if not self.start_secure_proxy(config):
                return 1
            
            # 启动 SSH 隧道
            if not self.start_tunnel(config):
                return 1
            
            print("\n🎉 Secure tunnel is running!")
            print("📊 Access logs will show below:")
            print("-" * 50)
            
            # 保持运行
            while True:
                time.sleep(1)
                
        except KeyboardInterrupt:
            print("\n🛑 Interrupted by user")
        finally:
            self.stop_all()
        
        return 0

def main():
    tunnel = SecureTunnel()
    return tunnel.run()

if __name__ == "__main__":
    sys.exit(main())