Upload 7 files

stream/__init__.py

@@ -0,0 +1,27 @@
+import asyncio
+import multiprocessing
+
+from stream import main
+
+def start(*args, **kwargs):
+    """
+    启动流式代理服务器，兼容位置参数和关键字参数
+
+    位置参数模式（与参考文件兼容）：
+        start(queue, port, proxy)
+
+    关键字参数模式：
+        start(queue=queue, port=port, proxy=proxy)
+    """
+    if args:
+        # 位置参数模式（与参考文件兼容）
+        queue = args[0] if len(args) > 0 else None
+        port = args[1] if len(args) > 1 else None
+        proxy = args[2] if len(args) > 2 else None
+    else:
+        # 关键字参数模式
+        queue = kwargs.get('queue', None)
+        port = kwargs.get('port', None)
+        proxy = kwargs.get('proxy', None)
+
+    asyncio.run(main.builtin(queue=queue, port=port, proxy=proxy))

stream/cert_manager.py

@@ -0,0 +1,171 @@
+import os
+import datetime
+from pathlib import Path
+from cryptography import x509
+from cryptography.x509.oid import NameOID
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.backends import default_backend
+
+class CertificateManager:
+    def __init__(self, cert_dir='certs'):
+        self.cert_dir = Path(cert_dir)
+        self.cert_dir.mkdir(exist_ok=True)
+        
+        self.ca_key_path = self.cert_dir / 'ca.key'
+        self.ca_cert_path = self.cert_dir / 'ca.crt'
+        
+        # Generate or load CA certificate
+        if not self.ca_cert_path.exists() or not self.ca_key_path.exists():
+            self._generate_ca_cert()
+        
+        self._load_ca_cert()
+    
+    def _generate_ca_cert(self):
+        """Generate a self-signed CA certificate"""
+        # Generate private key
+        private_key = rsa.generate_private_key(
+            public_exponent=65537,
+            key_size=2048,
+            backend=default_backend()
+        )
+        
+        # Write private key to file
+        with open(self.ca_key_path, 'wb') as f:
+            f.write(private_key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption()
+            ))
+        
+        # Create self-signed certificate
+        subject = issuer = x509.Name([
+            x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
+            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "California"),
+            x509.NameAttribute(NameOID.LOCALITY_NAME, "San Francisco"),
+            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Proxy CA"),
+            x509.NameAttribute(NameOID.COMMON_NAME, "Proxy CA Root"),
+        ])
+        
+        cert = x509.CertificateBuilder().subject_name(
+            subject
+        ).issuer_name(
+            issuer
+        ).public_key(
+            private_key.public_key()
+        ).serial_number(
+            x509.random_serial_number()
+        ).not_valid_before(
+            datetime.datetime.utcnow()
+        ).not_valid_after(
+            datetime.datetime.utcnow() + datetime.timedelta(days=3650)
+        ).add_extension(
+            x509.BasicConstraints(ca=True, path_length=None), critical=True
+        ).add_extension(
+            x509.KeyUsage(
+                digital_signature=True,
+                content_commitment=False,
+                key_encipherment=True,
+                data_encipherment=False,
+                key_agreement=False,
+                key_cert_sign=True,
+                crl_sign=True,
+                encipher_only=False,
+                decipher_only=False
+            ), critical=True
+        ).sign(private_key, hashes.SHA256(), default_backend())
+        
+        # Write certificate to file
+        with open(self.ca_cert_path, 'wb') as f:
+            f.write(cert.public_bytes(serialization.Encoding.PEM))
+    
+    def _load_ca_cert(self):
+        """Load the CA certificate and private key"""
+        with open(self.ca_key_path, 'rb') as f:
+            self.ca_key = serialization.load_pem_private_key(
+                f.read(),
+                password=None,
+                backend=default_backend()
+            )
+        
+        with open(self.ca_cert_path, 'rb') as f:
+            self.ca_cert = x509.load_pem_x509_certificate(
+                f.read(),
+                default_backend()
+            )
+    
+    def get_domain_cert(self, domain):
+        """Get or generate a certificate for the specified domain"""
+        cert_path = self.cert_dir / f"{domain}.crt"
+        key_path = self.cert_dir / f"{domain}.key"
+        
+        if cert_path.exists() and key_path.exists():
+            # Load existing certificate and key
+            with open(key_path, 'rb') as f:
+                private_key = serialization.load_pem_private_key(
+                    f.read(),
+                    password=None,
+                    backend=default_backend()
+                )
+            
+            with open(cert_path, 'rb') as f:
+                cert = x509.load_pem_x509_certificate(
+                    f.read(),
+                    default_backend()
+                )
+            
+            return private_key, cert
+        
+        # Generate new certificate
+        return self._generate_domain_cert(domain)
+    
+    def _generate_domain_cert(self, domain):
+        """Generate a certificate for the specified domain signed by the CA"""
+        # Generate private key
+        private_key = rsa.generate_private_key(
+            public_exponent=65537,
+            key_size=2048,
+            backend=default_backend()
+        )
+        
+        # Write private key to file
+        key_path = self.cert_dir / f"{domain}.key"
+        with open(key_path, 'wb') as f:
+            f.write(private_key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption()
+            ))
+        
+        # Create certificate
+        subject = x509.Name([
+            x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
+            x509.NameAttribute(NameOID.STATE_OR_PROVINCE_NAME, "California"),
+            x509.NameAttribute(NameOID.LOCALITY_NAME, "San Francisco"),
+            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Proxy Server"),
+            x509.NameAttribute(NameOID.COMMON_NAME, domain),
+        ])
+        
+        cert = x509.CertificateBuilder().subject_name(
+            subject
+        ).issuer_name(
+            self.ca_cert.subject
+        ).public_key(
+            private_key.public_key()
+        ).serial_number(
+            x509.random_serial_number()
+        ).not_valid_before(
+            datetime.datetime.utcnow()
+        ).not_valid_after(
+            datetime.datetime.utcnow() + datetime.timedelta(days=365)
+        ).add_extension(
+            x509.SubjectAlternativeName([x509.DNSName(domain)]),
+            critical=False
+        ).sign(self.ca_key, hashes.SHA256(), default_backend())
+        
+        # Write certificate to file
+        cert_path = self.cert_dir / f"{domain}.crt"
+        with open(cert_path, 'wb') as f:
+            f.write(cert.public_bytes(serialization.Encoding.PEM))
+        
+        return private_key, cert

stream/interceptors.py

@@ -0,0 +1,162 @@
+import json
+import logging
+import re
+import zlib
+
+class HttpInterceptor:
+    """
+    Class to intercept and process HTTP requests and responses
+    """
+    def __init__(self, log_dir='logs'):
+        self.log_dir = log_dir
+        self.logger = logging.getLogger('http_interceptor')
+        self.setup_logging()
+    
+    @staticmethod
+    def setup_logging():
+        """Set up logging configuration"""
+        logging.basicConfig(
+            level=logging.INFO,
+            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+            handlers=[
+                logging.StreamHandler()
+            ]
+        )
+    
+    @staticmethod
+    def should_intercept(host, path):
+        """
+        Determine if the request should be intercepted based on host and path
+        """
+        # Check if the endpoint contains GenerateContent
+        if 'GenerateContent' in path:
+            return True
+        
+        # Add more conditions as needed
+        return False
+    
+    async def process_request(self, request_data, host, path):
+        """
+        Process the request data before sending to the server
+        """
+        if not self.should_intercept(host, path):
+            return request_data
+        
+        # Log the request
+        self.logger.info(f"Intercepted request to {host}{path}")
+        
+        try:
+            return request_data
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            # Not JSON or not UTF-8, just pass through
+            return request_data
+    
+    async def process_response(self, response_data, host, path, headers):
+        """
+        Process the response data before sending to the client
+        """
+        try:
+            # Handle chunked encoding
+            decoded_data, is_done = self._decode_chunked(bytes(response_data))
+            # Handle gzip encoding
+            decoded_data = self._decompress_zlib_stream(decoded_data)
+            result  = self.parse_response(decoded_data)
+            result["done"] = is_done
+            return result
+        except Exception as e:
+            raise e
+
+    def parse_response(self, response_data):
+        pattern = rb'\[\[\[null,.*?]],"model"]'
+        matches = []
+        for match_obj in re.finditer(pattern, response_data):
+            matches.append(match_obj.group(0))
+
+
+        resp = {
+            "reason": "",
+            "body": "",
+            "function": [],
+        }
+
+        # Print each full match
+        for match in matches:
+            json_data = json.loads(match)
+
+            try:
+                payload = json_data[0][0]
+            except Exception as e:
+                continue
+
+            if len(payload)==2: # body
+                resp["body"] = resp["body"] + payload[1]
+            elif len(payload) == 11 and payload[1] is None and type(payload[10]) == list:  # function
+                array_tool_calls = payload[10]
+                func_name = array_tool_calls[0]
+                params = self.parse_toolcall_params(array_tool_calls[1])
+                resp["function"].append({"name":func_name, "params":params})
+            elif len(payload) > 2: # reason
+                resp["reason"] = resp["reason"] + payload[1]
+
+        return resp
+
+    def parse_toolcall_params(self, args):
+        try:
+            params = args[0]
+            func_params = {}
+            for param in params:
+                param_name = param[0]
+                param_value = param[1]
+
+                if type(param_value)==list:
+                    if len(param_value)==1: # null
+                        func_params[param_name] = None
+                    elif len(param_value) == 2: # number and integer
+                        func_params[param_name] = param_value[1]
+                    elif len(param_value) == 3: # string
+                        func_params[param_name] = param_value[2]
+                    elif len(param_value) == 4: # boolean
+                        func_params[param_name] = param_value[3] == 1
+                    elif len(param_value) == 5: # object
+                        func_params[param_name] = self.parse_toolcall_params(param_value[4])
+            return func_params
+        except Exception as e:
+            raise e
+
+    @staticmethod
+    def _decompress_zlib_stream(compressed_stream):
+        decompressor = zlib.decompressobj(wbits=zlib.MAX_WBITS | 32)  # zlib header
+        decompressed = decompressor.decompress(compressed_stream)
+        return decompressed
+
+    @staticmethod
+    def _decode_chunked(response_body: bytes) -> tuple[bytes, bool]:
+        chunked_data = bytearray()
+        while True:
+            # print(' '.join(format(x, '02x') for x in response_body))
+
+            length_crlf_idx = response_body.find(b"\r\n")
+            if length_crlf_idx == -1:
+                break
+
+            hex_length = response_body[:length_crlf_idx]
+            try:
+                length = int(hex_length, 16)
+            except ValueError as e:
+                logging.error(f"Parsing chunked length failed: {e}")
+                break
+
+            if length == 0:
+                length_crlf_idx = response_body.find(b"0\r\n\r\n")
+                if length_crlf_idx != -1:
+                    return chunked_data, True
+
+            if length + 2 > len(response_body):
+                break
+
+            chunked_data.extend(response_body[length_crlf_idx + 2:length_crlf_idx + 2 + length])
+            if length_crlf_idx + 2 + length + 2 > len(response_body):
+                break
+
+            response_body = response_body[length_crlf_idx + 2 + length + 2:]
+        return chunked_data, False

stream/main.py

@@ -0,0 +1,103 @@
+import argparse
+import asyncio
+import logging
+import multiprocessing
+import sys
+from pathlib import Path
+
+from stream.proxy_server import ProxyServer
+
+def parse_args():
+    """Parse command line arguments"""
+    parser = argparse.ArgumentParser(description='HTTPS Proxy Server with SSL Inspection')
+    
+    parser.add_argument('--host', default='127.0.0.1', help='Host to bind the proxy server')
+    parser.add_argument('--port', type=int, default=3120, help='Port to bind the proxy server')
+    parser.add_argument('--domains', nargs='+', default=['*.google.com'],
+                        help='List of domain patterns to intercept (regex)')
+    parser.add_argument('--proxy', help='Upstream proxy URL (e.g., http://user:pass@host:port)')
+    
+    return parser.parse_args()
+
+
+async def main():
+    """Main entry point"""
+    args = parse_args()
+
+    # Set up logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.StreamHandler()
+        ]
+    )
+    
+    logger = logging.getLogger('main')
+
+    # Create certs directory
+    cert_dir = Path('certs')
+    cert_dir.mkdir(exist_ok=True)
+    
+    # Print startup information
+    logger.info(f"Starting proxy server on {args.host}:{args.port}")
+    logger.info(f"Intercepting domains: {args.domains}")
+    if args.proxy:
+        logger.info(f"Using upstream proxy: {args.proxy}")
+    
+    # Create and start the proxy server
+    proxy_server = ProxyServer(
+        host=args.host,
+        port=args.port,
+        intercept_domains=args.domains,
+        upstream_proxy=args.proxy,
+        queue=None,
+    )
+
+    try:
+        await proxy_server.start()
+    except KeyboardInterrupt:
+        logger.info("Shutting down proxy server")
+    except Exception as e:
+        logger.error(f"Error starting proxy server: {e}")
+        sys.exit(1)
+
+
+async def builtin(queue: multiprocessing.Queue = None, port=None, proxy=None):
+    # Set up logging
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        handlers=[
+            logging.StreamHandler()
+        ]
+    )
+
+    logger = logging.getLogger('main')
+
+    # Create certs directory
+    cert_dir = Path('certs')
+    cert_dir.mkdir(exist_ok=True)
+
+    if port is None:
+        port = 3120
+
+    # Create and start the proxy server
+    proxy_server = ProxyServer(
+        host="127.0.0.1",
+        port=port,
+        intercept_domains=['*.google.com'],
+        upstream_proxy=proxy,
+        queue=queue,
+    )
+
+    try:
+        await proxy_server.start()
+    except KeyboardInterrupt:
+        logger.info("Shutting down proxy server")
+    except Exception as e:
+        logger.error(f"Error starting proxy server: {e}")
+        sys.exit(1)
+
+if __name__ == '__main__':
+    asyncio.run(main())

stream/proxy_connector.py

@@ -0,0 +1,68 @@
+import asyncio
+import ssl as ssl_module
+import urllib.parse
+from aiohttp import TCPConnector
+from python_socks.async_.asyncio import Proxy
+
+
+class ProxyConnector:
+    """
+    Class to handle connections through different types of proxies
+    """
+
+    def __init__(self, proxy_url=None):
+        self.proxy_url = proxy_url
+        self.connector = None
+
+        if proxy_url:
+            self._setup_connector()
+
+    def _setup_connector(self):
+        """Set up the appropriate connector based on the proxy URL"""
+        if not self.proxy_url:
+            self.connector = TCPConnector()
+            return
+
+        # Parse the proxy URL
+        parsed = urllib.parse.urlparse(self.proxy_url)
+        proxy_type = parsed.scheme.lower()
+
+        if proxy_type in ('http', 'https', 'socks4', 'socks5'):
+            self.connector = "SocksConnector"
+        else:
+            raise ValueError(f"Unsupported proxy type: {proxy_type}")
+
+    async def create_connection(self, host, port, ssl=None):
+        """Create a connection to the target host through the proxy"""
+        if not self.connector:
+            # Direct connection without proxy
+            reader, writer = await asyncio.open_connection(host, port, ssl=ssl)
+            return reader, writer
+
+        # SOCKS proxy connection
+        proxy = Proxy.from_url(self.proxy_url)
+        sock = await proxy.connect(dest_host=host, dest_port=port)
+        if ssl is None:
+            reader, writer = await asyncio.open_connection(
+                host=None,
+                port=None,
+                sock=sock,
+                ssl=None,
+            )
+            return reader, writer
+        else:
+            ssl_context = ssl_module.SSLContext(ssl_module.PROTOCOL_TLS_CLIENT)
+            ssl_context.check_hostname = False
+            ssl_context.verify_mode = ssl_module.CERT_NONE
+            ssl_context.minimum_version = ssl_module.TLSVersion.TLSv1_2  # Force TLS 1.2 or higher
+            ssl_context.maximum_version = ssl_module.TLSVersion.TLSv1_3  # Allow TLS 1.3 if supported
+            ssl_context.set_ciphers('DEFAULT@SECLEVEL=2')  # Use secure ciphers
+
+            reader, writer = await asyncio.open_connection(
+                host=None,
+                port=None,
+                sock=sock,
+                ssl=ssl_context,
+                server_hostname=host,
+            )
+            return reader, writer

stream/proxy_server.py

@@ -0,0 +1,357 @@
+import asyncio
+from typing import Optional
+import json
+import logging
+import ssl
+import multiprocessing
+from pathlib import Path
+from urllib.parse import urlparse
+
+from stream.cert_manager import CertificateManager
+from stream.proxy_connector import ProxyConnector
+from stream.interceptors import HttpInterceptor
+
+class ProxyServer:
+    """
+    Asynchronous HTTPS proxy server with SSL inspection capabilities
+    """
+    def __init__(self, host='0.0.0.0', port=3120, intercept_domains=None, upstream_proxy=None, queue: Optional[multiprocessing.Queue]=None):
+        self.host = host
+        self.port = port
+        self.intercept_domains = intercept_domains or []
+        self.upstream_proxy = upstream_proxy
+        self.queue = queue
+        
+        # Initialize components
+        self.cert_manager = CertificateManager()
+        self.proxy_connector = ProxyConnector(upstream_proxy)
+        
+        # Create logs directory
+        log_dir = Path('logs')
+        log_dir.mkdir(exist_ok=True)
+        self.interceptor = HttpInterceptor(str(log_dir))
+        
+        # Set up logging
+        self.logger = logging.getLogger('proxy_server')
+    
+    def should_intercept(self, host):
+        """
+        Determine if the connection to the host should be intercepted
+        """
+        if host in self.intercept_domains:
+            return True
+
+        # Wildcard match (e.g. *.example.com)
+        for d in self.intercept_domains:
+            if d.startswith("*."):
+                suffix = d[1:]  # Remove *
+                if host.endswith(suffix):
+                    return True
+
+        return False
+
+    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
+        """
+        Handle a client connection
+        """
+        try:
+            # Read the initial request line
+            request_line = await reader.readline()
+            request_line = request_line.decode('utf-8').strip()
+            
+            if not request_line:
+                writer.close()
+                return
+            
+            # Parse the request line
+            method, target, version = request_line.split(' ')
+            
+            if method == 'CONNECT':
+                # Handle HTTPS connection
+                await self._handle_connect(reader, writer, target)
+
+        except Exception as e:
+            self.logger.error(f"Error handling client: {e}")
+        finally:
+            writer.close()
+    
+    async def _handle_connect(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter, target: str):
+        """
+        Handle CONNECT method (for HTTPS connections)
+        """
+
+        host, port = target.split(':')
+        port = int(port)
+        # Determine if we should intercept this connection
+        intercept = self.should_intercept(host)
+
+        if intercept:
+            self.logger.info(f"Sniff HTTPS requests to : {target}")
+
+            self.cert_manager.get_domain_cert(host)
+
+            # Send 200 Connection Established to the client
+            writer.write(b'HTTP/1.1 200 Connection Established\r\n\r\n')
+            await writer.drain()
+
+            # Drop the proxy connect header
+            await reader.read(8192)
+
+            loop = asyncio.get_running_loop()
+            transport = writer.transport # This is the original client transport
+
+            if transport is None: # 新增检查块开始
+                self.logger.warning(f"Client writer transport is None for {host}:{port} before TLS upgrade. Closing.")
+                # writer is likely already closed or in a bad state.
+                # We can't proceed with start_tls if transport is None.
+                return # Exit _handle_connect for this client # 新增检查块结束
+
+            ssl_context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+            ssl_context.load_cert_chain(
+                certfile=self.cert_manager.cert_dir / f"{host}.crt",
+                keyfile=self.cert_manager.cert_dir / f"{host}.key"
+            )
+
+            # 1. 正确获取与原始 transport 关联的协议实例
+            # 'transport' here is 'writer.transport' from line 101, now checked not to be None
+            client_protocol = transport.get_protocol()
+
+            # 2. 将获取到的 client_protocol 实例传递给 start_tls
+            #    loop.start_tls 会修改这个 client_protocol 实例，使其与 new_transport 关联
+            new_transport = await loop.start_tls(
+                transport=transport,
+                protocol=client_protocol,  # 关键：传递获取到的协议实例
+                sslcontext=ssl_context,
+                server_side=True
+            )
+
+            # 3. 增加对 new_transport 的 None 检查 (主要为了类型安全和 Pylance)
+            if new_transport is None:
+                self.logger.error(f"loop.start_tls returned None for {host}:{port}, which is unexpected. Closing connection.")
+                # Ensure client writer is closed if it was opened or transport was valid before
+                writer.close()
+                # await writer.wait_closed() # Consider if waiting is necessary here
+                return
+            
+            client_reader = reader
+
+            # 4. 创建 StreamWriter 时，使用被 start_tls 正确更新过的 client_protocol
+            client_writer = asyncio.StreamWriter(
+                transport=new_transport,    # 使用新的 TLS transport
+                protocol=client_protocol,   # 关键：使用被 start_tls 更新过的协议实例
+                reader=client_reader,
+                loop=loop
+            )
+
+            # Connect to the target server
+            try:
+                server_reader, server_writer = await self.proxy_connector.create_connection(
+                    host, port, ssl=ssl.create_default_context()
+                )
+                
+                # Start bidirectional forwarding with interception
+                await self._forward_data_with_interception(
+                    client_reader, client_writer,
+                    server_reader, server_writer,
+                    host
+                )
+            except Exception as e:
+                # self.logger.error(f"Error connecting to server {host}:{port}: {e}")
+                client_writer.close()
+                # await client_writer.wait_closed()
+        else:
+            # No interception, just forward the connection
+            writer.write(b'HTTP/1.1 200 Connection Established\r\n\r\n')
+            await writer.drain()
+
+            # Drop the proxy connect header
+            await reader.read(8192)
+
+            try:
+                # Connect to the target server
+                server_reader, server_writer = await self.proxy_connector.create_connection(
+                    host, port, ssl=None
+                )
+
+                # Start bidirectional forwarding without interception
+                await self._forward_data(
+                    reader, writer,
+                    server_reader, server_writer
+                )
+            except Exception as e:
+                # self.logger.error(f"Error connecting to server {host}:{port}: {e}")
+                writer.close()
+                # await writer.wait_closed()
+    async def _forward_data(self, client_reader, client_writer, server_reader, server_writer):
+        """
+        Forward data between client and server without interception
+        """
+        async def _forward(reader, writer):
+            try:
+                while True:
+                    data = await reader.read(8192)
+                    if not data:
+                        break
+                    writer.write(data)
+                    await writer.drain()
+            except Exception as e:
+                self.logger.error(f"Error forwarding data: {e}")
+            finally:
+                writer.close()
+        
+        # Create tasks for both directions
+        client_to_server = asyncio.create_task(_forward(client_reader, server_writer))
+        server_to_client = asyncio.create_task(_forward(server_reader, client_writer))
+        
+        # Wait for both tasks to complete
+        tasks = [client_to_server, server_to_client]
+        await asyncio.gather(*tasks)
+        # await asyncio.gather(client_to_server, server_to_client)
+    
+    async def _forward_data_with_interception(self, client_reader, client_writer, 
+                                             server_reader, server_writer, host):
+        """
+        Forward data between client and server with interception
+        """
+        # Buffer to store HTTP request/response data
+        client_buffer = bytearray()
+        server_buffer = bytearray()
+        should_sniff = False
+
+        # Parse HTTP headers from client
+        async def _process_client_data():
+            nonlocal client_buffer, should_sniff
+            
+            try:
+                while True:
+                    data = await client_reader.read(8192)
+                    if not data:
+                        break
+                    client_buffer.extend(data)
+                    
+                    # Try to parse HTTP request
+                    if b'\r\n\r\n' in client_buffer:
+                        # Split headers and body
+                        headers_end = client_buffer.find(b'\r\n\r\n') + 4
+                        headers_data = client_buffer[:headers_end]
+                        body_data = client_buffer[headers_end:]
+                        
+                        # Parse request line and headers
+                        lines = headers_data.split(b'\r\n')
+                        request_line = lines[0].decode('utf-8')
+                        
+                        try:
+                            method, path, _ = request_line.split(' ')
+                        except ValueError:
+                            # Not a valid HTTP request, just forward
+                            server_writer.write(client_buffer)
+                            await server_writer.drain()
+                            client_buffer.clear()
+                            continue
+                        
+                        # Check if we should intercept this request
+                        if 'GenerateContent' in path:
+                            should_sniff = True
+                            # Process the request body
+                            processed_body = await self.interceptor.process_request(
+                                body_data, host, path
+                            )
+                            
+                            # Send the processed request
+                            server_writer.write(headers_data)
+                            server_writer.write(processed_body)
+                        else:
+                            should_sniff = False
+                            # Forward the request as is
+                            server_writer.write(client_buffer)
+                        
+                        await server_writer.drain()
+                        client_buffer.clear()
+                    else:
+                        # Not enough data to parse headers, forward as is
+                        server_writer.write(data)
+                        await server_writer.drain()
+                        client_buffer.clear()
+            except Exception as e:
+                self.logger.error(f"Error processing client data: {e}")
+            finally:
+                server_writer.close()
+                # await server_writer.wait_closed()
+        
+        # Parse HTTP headers from server
+        async def _process_server_data():
+            nonlocal server_buffer, should_sniff
+            
+            try:
+                while True:
+                    data = await server_reader.read(8192)
+                    if not data:
+                        break
+
+                    server_buffer.extend(data)
+                    if b'\r\n\r\n' in server_buffer:
+                        # Split headers and body
+                        headers_end = server_buffer.find(b'\r\n\r\n') + 4
+                        headers_data = server_buffer[:headers_end]
+                        body_data = server_buffer[headers_end:]
+
+                        # Parse status line and headers
+                        lines = headers_data.split(b'\r\n')
+
+                        # Parse headers
+                        headers = {}
+                        for i in range(1, len(lines)):
+                            if not lines[i]:
+                                continue
+                            try:
+                                key, value = lines[i].decode('utf-8').split(':', 1)
+                                headers[key.strip()] = value.strip()
+                            except ValueError:
+                                continue
+
+                        # Check if this is a response to a GenerateContent request
+                        if should_sniff:
+                            try:
+                                resp = await self.interceptor.process_response(
+                                    body_data, host, "", headers
+                                )
+
+                                if self.queue is not None:
+                                    self.queue.put(json.dumps(resp))
+                            except Exception as e:
+                                pass
+
+                    # Not enough data to parse headers, forward as is
+                    client_writer.write(data)
+                    # await client_writer.drain()
+                    if b"0\r\n\r\n" in server_buffer:
+                        server_buffer.clear()
+            except Exception as e:
+                self.logger.error(f"Error processing server data: {e}")
+            finally:
+                client_writer.close()
+                # await client_writer.wait_closed()
+        
+        # Create tasks for both directions
+        client_to_server = asyncio.create_task(_process_client_data())
+        server_to_client = asyncio.create_task(_process_server_data())
+
+
+        # Wait for both tasks to complete
+        tasks = [client_to_server, server_to_client]
+        await asyncio.gather(*tasks)
+        # await asyncio.gather(client_to_server, server_to_client)
+    
+    async def start(self):
+        """
+        Start the proxy server
+        """
+        server = await asyncio.start_server(
+            self.handle_client, self.host, self.port
+        )
+        
+        addr = server.sockets[0].getsockname()
+        self.logger.info(f'Serving on {addr}')
+        
+        async with server:
+            await server.serve_forever()

stream/utils.py

@@ -0,0 +1,58 @@
+import logging
+from urllib.parse import urlparse
+
+def is_generate_content_endpoint(url):
+    """
+    Check if the URL is a GenerateContent endpoint
+    """
+    return 'GenerateContent' in url
+
+def parse_proxy_url(proxy_url):
+    """
+    Parse a proxy URL into its components
+    
+    Returns:
+        tuple: (scheme, host, port, username, password)
+    """
+    if not proxy_url:
+        return None, None, None, None, None
+    
+    parsed = urlparse(proxy_url)
+    
+    scheme = parsed.scheme
+    host = parsed.hostname
+    port = parsed.port
+    username = parsed.username
+    password = parsed.password
+    
+    return scheme, host, port, username, password
+
+def setup_logger(name, log_file=None, level=logging.INFO):
+    """
+    Set up a logger with the specified name and configuration
+    
+    Args:
+        name (str): Logger name
+        log_file (str, optional): Path to log file
+        level (int, optional): Logging level
+        
+    Returns:
+        logging.Logger: Configured logger
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+    
+    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    
+    # Add console handler
+    console_handler = logging.StreamHandler()
+    console_handler.setFormatter(formatter)
+    logger.addHandler(console_handler)
+    
+    # Add file handler if specified
+    if log_file:
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+    
+    return logger