112 lines
4.0 KiB
Python
112 lines
4.0 KiB
Python
import sys
|
|
import argparse
|
|
import platform
|
|
import ctypes
|
|
import winreg
|
|
from mitmproxy.tools.main import mitmdump
|
|
from proxy_addon import load_config
|
|
|
|
|
|
class SystemProxyManager:
|
|
def __init__(self, host: str, port: int):
|
|
self.host = host
|
|
self.port = port
|
|
self.original_enable = None
|
|
self.original_server = None
|
|
|
|
def _apply_windows_internet_options(self):
|
|
option_refresh = 37
|
|
option_settings_changed = 39
|
|
internet_set_option = ctypes.windll.Wininet.InternetSetOptionW
|
|
internet_set_option(0, option_settings_changed, 0, 0)
|
|
internet_set_option(0, option_refresh, 0, 0)
|
|
|
|
def enable(self):
|
|
if platform.system().lower() != "windows":
|
|
return
|
|
key_path = r"Software\Microsoft\Windows\CurrentVersion\Internet Settings"
|
|
key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, key_path, 0, winreg.KEY_READ | winreg.KEY_WRITE)
|
|
try:
|
|
self.original_enable, _ = winreg.QueryValueEx(key, "ProxyEnable")
|
|
except FileNotFoundError:
|
|
self.original_enable = 0
|
|
try:
|
|
self.original_server, _ = winreg.QueryValueEx(key, "ProxyServer")
|
|
except FileNotFoundError:
|
|
self.original_server = ""
|
|
winreg.SetValueEx(key, "ProxyEnable", 0, winreg.REG_DWORD, 1)
|
|
winreg.SetValueEx(key, "ProxyServer", 0, winreg.REG_SZ, f"{self.host}:{self.port}")
|
|
winreg.CloseKey(key)
|
|
self._apply_windows_internet_options()
|
|
|
|
def disable(self):
|
|
if platform.system().lower() != "windows":
|
|
return
|
|
if self.original_enable is None or self.original_server is None:
|
|
return
|
|
key_path = r"Software\Microsoft\Windows\CurrentVersion\Internet Settings"
|
|
key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, key_path, 0, winreg.KEY_READ | winreg.KEY_WRITE)
|
|
winreg.SetValueEx(key, "ProxyEnable", 0, winreg.REG_DWORD, self.original_enable)
|
|
winreg.SetValueEx(key, "ProxyServer", 0, winreg.REG_SZ, self.original_server)
|
|
winreg.CloseKey(key)
|
|
self._apply_windows_internet_options()
|
|
|
|
|
|
def start_proxy(config_path: str = "config.json", port: int = 8080, host: str = "127.0.0.1", enable_system_proxy: bool = True):
|
|
config = load_config(config_path)
|
|
|
|
proxy_config = config.get('proxy', {})
|
|
listen_port = port or proxy_config.get('listen_port', 8080)
|
|
listen_host = host or proxy_config.get('listen_host', '127.0.0.1')
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"LLM Proxy Server")
|
|
print(f"{'='*60}")
|
|
print(f"Listening on: {listen_host}:{listen_port}")
|
|
print(f"Config file: {config_path}")
|
|
print(f"Database: {config.get('database', {}).get('path', 'llm_data.db')}")
|
|
if enable_system_proxy and platform.system().lower() == "windows":
|
|
print("System proxy: enabled for current session")
|
|
print(f"{'='*60}\n")
|
|
|
|
manager = None
|
|
if enable_system_proxy:
|
|
manager = SystemProxyManager(listen_host, listen_port)
|
|
manager.enable()
|
|
|
|
sys.argv = [
|
|
'mitmdump',
|
|
'-q',
|
|
'-s', 'proxy_addon.py',
|
|
'--listen-host', listen_host,
|
|
'--listen-port', str(listen_port),
|
|
'--set', 'block_global=false',
|
|
'--set', 'flow_detail=0'
|
|
]
|
|
|
|
try:
|
|
mitmdump()
|
|
finally:
|
|
if manager is not None:
|
|
manager.disable()
|
|
|
|
|
|
def cli_main():
|
|
parser = argparse.ArgumentParser(description='Start LLM Proxy Server')
|
|
parser.add_argument('--config', '-c', type=str, default='config.json',
|
|
help='Path to config file')
|
|
parser.add_argument('--port', '-p', type=int, default=None,
|
|
help='Listen port (overrides config)')
|
|
parser.add_argument('--host', '-H', type=str, default=None,
|
|
help='Listen host (overrides config)')
|
|
parser.add_argument('--no-system-proxy', action='store_true',
|
|
help='Do not modify system proxy settings')
|
|
|
|
args = parser.parse_args()
|
|
|
|
start_proxy(args.config, args.port, args.host, not args.no_system_proxy)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
cli_main()
|