diff --git a/.gitignore b/.gitignore index 3a2b94ee..df3855a4 100644 --- a/.gitignore +++ b/.gitignore @@ -45,3 +45,6 @@ weaviate_data/ # Codegate Dashboard DB codegate.db + +# certificate directory +certs/ diff --git a/cert_gen.py b/cert_gen.py new file mode 100644 index 00000000..d797bc62 --- /dev/null +++ b/cert_gen.py @@ -0,0 +1,160 @@ +from cryptography import x509 +from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +import datetime +import os + +def generate_certificates(cert_dir="certs"): + """Generate self-signed certificates with proper extensions for HTTPS proxy""" + # Create certificates directory if it doesn't exist + if not os.path.exists(cert_dir): + print("Making: ", cert_dir) + os.makedirs(cert_dir) + + # Generate private key + ca_private_key = rsa.generate_private_key( + public_exponent=65537, + key_size=4096, # Increased key size for better security + ) + + # Generate public key + ca_public_key = ca_private_key.public_key() + + # CA BEGIN + name = x509.Name([ + x509.NameAttribute(NameOID.COMMON_NAME, "Proxy Pilot CA"), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Proxy Pilot"), + x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Development"), + x509.NameAttribute(NameOID.COUNTRY_NAME, "UK"), + ]) + + builder = x509.CertificateBuilder() + builder = builder.subject_name(name) + builder = builder.issuer_name(name) + builder = builder.public_key(ca_public_key) + builder = builder.serial_number(x509.random_serial_number()) + builder = builder.not_valid_before(datetime.datetime.utcnow()) + builder = builder.not_valid_after( + datetime.datetime.utcnow() + datetime.timedelta(days=3650) # 10 years + ) + + builder = builder.add_extension( + x509.BasicConstraints(ca=True, path_length=None), + critical=True, + ) + + builder = builder.add_extension( + x509.KeyUsage( + digital_signature=True, + content_commitment=False, + key_encipherment=True, + data_encipherment=False, + key_agreement=False, + key_cert_sign=True, # This is a CA + crl_sign=True, + encipher_only=False, + decipher_only=False + ), + critical=True, + ) + + ca_cert = builder.sign( + private_key=ca_private_key, + algorithm=hashes.SHA256(), + ) + + # Save CA certificate and key + + with open("certs/ca.crt", "wb") as f: + f.write(ca_cert.public_bytes(serialization.Encoding.PEM)) + + with open("certs/ca.key", "wb") as f: + f.write(ca_private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption() + )) + # CA END + + # SERVER BEGIN + + ## Generate new certificate for domain + server_key = rsa.generate_private_key( + public_exponent=65537, + key_size=2048, # 2048 bits is sufficient for domain certs + ) + + name = x509.Name([ + x509.NameAttribute(NameOID.COMMON_NAME, "Proxy Pilot CA"), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Proxy Pilot Generated"), + ]) + + builder = x509.CertificateBuilder() + builder = builder.subject_name(name) + builder = builder.issuer_name(ca_cert.subject) + builder = builder.public_key(server_key.public_key()) + builder = builder.serial_number(x509.random_serial_number()) + builder = builder.not_valid_before(datetime.datetime.utcnow()) + builder = builder.not_valid_after( + datetime.datetime.utcnow() + datetime.timedelta(days=365) + ) + + # Add domain to SAN + builder = builder.add_extension( + x509.SubjectAlternativeName([x509.DNSName("localhost")]), + critical=False, + ) + + # Add extended key usage + builder = builder.add_extension( + x509.ExtendedKeyUsage([ + ExtendedKeyUsageOID.SERVER_AUTH, + ExtendedKeyUsageOID.CLIENT_AUTH, + ]), + critical=False, + ) + + # Basic constraints (not a CA) + builder = builder.add_extension( + x509.BasicConstraints(ca=False, path_length=None), + critical=True, + ) + + certificate = builder.sign( + private_key=ca_private_key, + algorithm=hashes.SHA256(), + ) + + with open("certs/server.crt", "wb") as f: + f.write(certificate.public_bytes(serialization.Encoding.PEM)) + + with open("certs/server.key", "wb") as f: + f.write(server_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption() + )) + + + print("Certificates generated successfully in the 'certs' directory") + print("\nTo trust these certificates:") + print("\nOn macOS:") + print("sudo security add-trusted-cert -d -r trustRoot -k /Library/Keychains/System.keychain certs/server.crt") + print("\nOn Windows (PowerShell as Admin):") + print("Import-Certificate -FilePath \"certs\\server.crt\" -CertStoreLocation Cert:\\LocalMachine\\Root") + print("\nOn Linux:") + print("sudo cp certs/server.crt /usr/local/share/ca-certificates/proxy-pilot.crt") + print("sudo update-ca-certificates") + print("\nFor VSCode, add to settings.json:") + print('''{ + "http.proxy": "https://localhost:8989", + "http.proxySupport": "on", + "github.copilot.advanced": { + "debug.testOverrideProxyUrl": "https://localhost:8989", + "debug.overrideProxyUrl": "https://localhost:8989" + } +}''') + +if __name__ == "__main__": + generate_certificates() diff --git a/config.yaml.example b/config.yaml.example index 5d06f006..28d3171a 100644 --- a/config.yaml.example +++ b/config.yaml.example @@ -2,14 +2,15 @@ # Network settings port: 8989 # Port to listen on (1-65535) +proxy_port: 8990 # Proxy port to listen on (1-65535) host: "localhost" # Host to bind to (use localhost for all interfaces) # Logging configuration log_level: "INFO" # One of: ERROR, WARNING, INFO, DEBUG # Note: This configuration can be overridden by: -# 1. CLI arguments (--port, --host, --log-level) -# 2. Environment variables (CODEGATE_APP_PORT, CODEGATE_APP_HOST, CODEGATE_APP_LOG_LEVEL) +# 1. CLI arguments (--port, --proxy-port, --host, --log-level) +# 2. Environment variables (CODEGATE_APP_PORT, CODEGATE_APP_PROXY_PORT, CODEGATE_APP_HOST, CODEGATE_APP_LOG_LEVEL) # Provider URLs provider_urls: @@ -26,6 +27,22 @@ provider_urls: # --openai-url # --anthropic-url +# Certificate configuration +certs_dir: "./certs" # Directory for certificate files +ca_cert: "ca.crt" # CA certificate file name +ca_key: "ca.key" # CA key file name +server_cert: "server.crt" # Server certificate file name +server_key: "server.key" # Server key file name + +# Note: Certificate configuration can be overridden by: +# 1. CLI arguments (--certs-dir, --ca-cert, --ca-key, --server-cert, --server-key) +# 2. Environment variables: +# CODEGATE_CERTS_DIR +# CODEGATE_CA_CERT +# CODEGATE_CA_KEY +# CODEGATE_SERVER_CERT +# CODEGATE_SERVER_KEY + # Embedding model configuration #### diff --git a/docs/configuration.md b/docs/configuration.md index e4971fdc..9daf66de 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -12,6 +12,7 @@ The configuration system in Codegate is managed through the `Config` class in `c ## Default Configuration Values - Port: 8989 +- Proxy Port: 8990 - Host: "localhost" - Log Level: "INFO" - Log Format: "JSON" @@ -21,6 +22,12 @@ The configuration system in Codegate is managed through the `Config` class in `c - OpenAI: "https://api.openai.com/v1" - Anthropic: "https://api.anthropic.com/v1" - Ollama: "http://localhost:11434" +- Certificate Configuration: + - Certs Directory: "./certs" + - CA Certificate: "ca.crt" + - CA Key: "ca.key" + - Server Certificate: "server.crt" + - Server Key: "server.key" ## Configuration Methods @@ -35,6 +42,7 @@ config = Config.from_file("config.yaml") Example config.yaml: ```yaml port: 8989 +proxy_port: 8990 host: localhost log_level: INFO log_format: JSON @@ -43,6 +51,11 @@ provider_urls: openai: "https://api.openai.com/v1" anthropic: "https://api.anthropic.com/v1" ollama: "http://localhost:11434" +certs_dir: "./certs" +ca_cert: "ca.crt" +ca_key: "ca.key" +server_cert: "server.crt" +server_key: "server.key" ``` ### From Environment Variables @@ -50,6 +63,7 @@ provider_urls: Environment variables are automatically loaded with these mappings: - `CODEGATE_APP_PORT`: Server port +- `CODEGATE_APP_PROXY_PORT`: Server proxy port - `CODEGATE_APP_HOST`: Server host - `CODEGATE_APP_LOG_LEVEL`: Logging level - `CODEGATE_LOG_FORMAT`: Log format @@ -58,6 +72,11 @@ Environment variables are automatically loaded with these mappings: - `CODEGATE_PROVIDER_OPENAI_URL`: OpenAI provider URL - `CODEGATE_PROVIDER_ANTHROPIC_URL`: Anthropic provider URL - `CODEGATE_PROVIDER_OLLAMA_URL`: Ollama provider URL +- `CODEGATE_CERTS_DIR`: Directory for certificate files +- `CODEGATE_CA_CERT`: CA certificate file name +- `CODEGATE_CA_KEY`: CA key file name +- `CODEGATE_SERVER_CERT`: Server certificate file name +- `CODEGATE_SERVER_KEY`: Server key file name ```python config = Config.from_env() @@ -65,6 +84,29 @@ config = Config.from_env() ## Configuration Options +### Network Settings + +Network settings can be configured in several ways: + +1. In Configuration File: + ```yaml + port: 8989 # Port to listen on (1-65535) + proxy_port: 8990 # Proxy port to listen on (1-65535) + host: "localhost" # Host to bind to + ``` + +2. Via Environment Variables: + ```bash + export CODEGATE_APP_PORT=8989 + export CODEGATE_APP_PROXY_PORT=8990 + export CODEGATE_APP_HOST=localhost + ``` + +3. Via CLI Flags: + ```bash + codegate serve --port 8989 --proxy-port 8990 --host localhost + ``` + ### Provider URLs Provider URLs can be configured in several ways: @@ -95,6 +137,33 @@ Note: - For the vLLM provider, the /v1 path is automatically appended to the base URL if not present. - For the Ollama provider, the /api path is automatically appended to the base URL if not present. +### Certificate Configuration + +Certificate files can be configured in several ways: + +1. In Configuration File: + ```yaml + certs_dir: "./certs" + ca_cert: "ca.crt" + ca_key: "ca.key" + server_cert: "server.crt" + server_key: "server.key" + ``` + +2. Via Environment Variables: + ```bash + export CODEGATE_CERTS_DIR=./certs + export CODEGATE_CA_CERT=ca.crt + export CODEGATE_CA_KEY=ca.key + export CODEGATE_SERVER_CERT=server.crt + export CODEGATE_SERVER_KEY=server.key + ``` + +3. Via CLI Flags: + ```bash + codegate serve --certs-dir ./certs --ca-cert ca.crt --ca-key ca.key --server-cert server.crt --server-key server.key + ``` + ### Log Levels Available log levels (case-insensitive): @@ -160,6 +229,7 @@ prompt = config.prompts.prompt_name The configuration system uses a custom `ConfigurationError` exception for handling configuration-related errors, such as: - Invalid port numbers (must be between 1 and 65535) +- Invalid proxy port numbers (must be between 1 and 65535) - Invalid log levels - Invalid log formats - YAML parsing errors diff --git a/src/codegate/ca/codegate_ca.py b/src/codegate/ca/codegate_ca.py new file mode 100644 index 00000000..71587cc2 --- /dev/null +++ b/src/codegate/ca/codegate_ca.py @@ -0,0 +1,404 @@ +from cryptography import x509 +from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +import datetime +import os +import ssl +import structlog +from typing import Tuple, Dict, Optional +from codegate.config import Config + +logger = structlog.get_logger("codegate") + +class CertificateAuthority: + """ + Singleton class for Certificate Authority management. + Access the instance using CertificateAuthority.get_instance() + """ + _instance: Optional['CertificateAuthority'] = None + + @classmethod + def get_instance(cls) -> 'CertificateAuthority': + """Get or create the singleton instance of CertificateAuthority""" + if cls._instance is None: + logger.debug("Creating new CertificateAuthority instance") + cls._instance = cls() + return cls._instance + + def __init__(self): + """ + Initialize the Certificate Authority. + Note: Use get_instance() instead of creating a new instance directly. + """ + if CertificateAuthority._instance is not None: + raise RuntimeError("Use CertificateAuthority.get_instance() instead") + + logger.debug("Initializing Certificate Authority class: CertificateAuthority") + self._ca_cert = None + self._ca_key = None + self._cert_cache: Dict[str, Tuple[str, str]] = {} + self._load_or_generate_ca() + CertificateAuthority._instance = self + + def _load_or_generate_ca(self): + """Load existing CA certificate and key or generate new ones""" + logger.debug("Loading or generating CA certificate and key: fn: _load_or_generate_ca") + ca_cert = os.path.join(Config.get_config().certs_dir, Config.get_config().ca_cert) + ca_key = os.path.join(Config.get_config().certs_dir, Config.get_config().ca_key) + + if os.path.exists(ca_cert) and os.path.exists(ca_key): + # Load existing CA certificate and key + with open(ca_cert, "rb") as f: + logger.debug(f"Loading CA certificate from {ca_cert}") + self._ca_cert = x509.load_pem_x509_certificate(f.read()) + with open(ca_key, "rb") as f: + logger.debug(f"Loading CA key from {ca_key}") + self._ca_key = serialization.load_pem_private_key(f.read(), password=None) + else: + # Generate new CA certificate and key + logger.debug("Generating new CA certificate and key") + self._ca_key = rsa.generate_private_key( + public_exponent=65537, + key_size=4096, + ) + + name = x509.Name([ + x509.NameAttribute(NameOID.COMMON_NAME, "CodeGate CA"), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "CodeGate"), + x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "CodeGate"), + x509.NameAttribute(NameOID.COUNTRY_NAME, "UK"), + ]) + + builder = x509.CertificateBuilder() + builder = builder.subject_name(name) + builder = builder.issuer_name(name) + builder = builder.public_key(self._ca_key.public_key()) + builder = builder.serial_number(x509.random_serial_number()) + builder = builder.not_valid_before(datetime.datetime.utcnow()) + builder = builder.not_valid_after( + datetime.datetime.utcnow() + datetime.timedelta(days=3650) # 10 years + ) + + builder = builder.add_extension( + x509.BasicConstraints(ca=True, path_length=None), + critical=True, + ) + + builder = builder.add_extension( + x509.KeyUsage( + digital_signature=True, + content_commitment=False, + key_encipherment=True, + data_encipherment=False, + key_agreement=False, + key_cert_sign=True, # This is a CA + crl_sign=True, + encipher_only=False, + decipher_only=False + ), + critical=True, + ) + + self._ca_cert = builder.sign( + private_key=self._ca_key, + algorithm=hashes.SHA256(), + ) + + # Save CA certificate and key + if not os.path.exists(Config.get_config().certs_dir): + logger.debug(f"Creating directory: {Config.get_config().certs_dir}") + os.makedirs(Config.get_config().certs_dir) + + with open(ca_cert, "wb") as f: + logger.debug(f"Saving CA certificate to {ca_cert}") + f.write(self._ca_cert.public_bytes(serialization.Encoding.PEM)) + + with open(ca_key, "wb") as f: + logger.debug(f"Saving CA key to {ca_key}") + f.write(self._ca_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption() + )) + + def get_domain_certificate(self, domain: str) -> Tuple[str, str]: + """Get or generate a certificate for a specific domain""" + logger.debug(f"Getting domain certificate for domain: {domain} fn: get_domain_certificate") + if domain in self._cert_cache: + return self._cert_cache[domain] + + # Generate new certificate for domain + logger.debug(f"Generating private key for domain: {domain}") + key = rsa.generate_private_key( + public_exponent=65537, + key_size=2048, # 2048 bits is sufficient for domain certs + ) + + name = x509.Name([ + x509.NameAttribute(NameOID.COMMON_NAME, domain), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Proxy Pilot Generated"), + ]) + + builder = x509.CertificateBuilder() + builder = builder.subject_name(name) + builder = builder.issuer_name(self._ca_cert.subject) + builder = builder.public_key(key.public_key()) + builder = builder.serial_number(x509.random_serial_number()) + builder = builder.not_valid_before(datetime.datetime.utcnow()) + builder = builder.not_valid_after( + datetime.datetime.utcnow() + datetime.timedelta(days=365) + ) + + # Add domain to SAN + builder = builder.add_extension( + x509.SubjectAlternativeName([x509.DNSName(domain)]), + critical=False, + ) + + # Add extended key usage + builder = builder.add_extension( + x509.ExtendedKeyUsage([ + ExtendedKeyUsageOID.SERVER_AUTH, + ExtendedKeyUsageOID.CLIENT_AUTH, + ]), + critical=False, + ) + + # Basic constraints (not a CA) + builder = builder.add_extension( + x509.BasicConstraints(ca=False, path_length=None), + critical=True, + ) + + logger.debug(f"Signing certificate for domain: {domain}") + certificate = builder.sign( + private_key=self._ca_key, + algorithm=hashes.SHA256(), + ) + + # Save domain certificate and key + logger.debug(f"Saving certificate and key for domain: {domain}") + domain_cert_path = os.path.join(Config.get_config().certs_dir, f"{domain}.crt") + domain_key_path = os.path.join(Config.get_config().certs_dir, f"{domain}.key") + + with open(domain_cert_path, "wb") as f: + logger.debug(f"Saving certificate to {domain_cert_path}") + f.write(certificate.public_bytes(serialization.Encoding.PEM)) + + with open(domain_key_path, "wb") as f: + logger.debug(f"Saving key to {domain_key_path}") + f.write(key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption() + )) + + self._cert_cache[domain] = (domain_cert_path, domain_key_path) + return domain_cert_path, domain_key_path + + def generate_certificates(self) -> Tuple[str, str]: + """Generate self-signed certificates with proper extensions for HTTPS proxy""" + logger.debug("Generating certificates fn: generate_certificates") + + if not os.path.exists(Config.get_config().certs_dir): + logger.debug(f"Creating directory: {Config.get_config().certs_dir}") + os.makedirs(Config.get_config().certs_dir) + + # Generate private key + logger.debug("Generating private key for CA") + ca_private_key = rsa.generate_private_key( + public_exponent=65537, + key_size=4096, + ) + + # Generate public key + logger.debug("Generating public key for CA") + ca_public_key = ca_private_key.public_key() + + # Add name attributes + name = x509.Name([ + x509.NameAttribute(NameOID.COMMON_NAME, "CodeGate CA"), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "CodeGate"), + x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "CodeGate"), + x509.NameAttribute(NameOID.COUNTRY_NAME, "UK"), + ]) + + # Create certificate builder + builder = x509.CertificateBuilder() + + # Basic certificate information + builder = builder.subject_name(name) + builder = builder.issuer_name(name) + builder = builder.public_key(ca_public_key) + builder = builder.serial_number(x509.random_serial_number()) + builder = builder.not_valid_before(datetime.datetime.utcnow()) + builder = builder.not_valid_after( + datetime.datetime.utcnow() + datetime.timedelta(days=365) + ) + + builder = builder.add_extension( + x509.BasicConstraints(ca=True, path_length=None), + critical=True, + ) + + # Add key usage extensions + builder = builder.add_extension( + x509.KeyUsage( + digital_signature=True, + content_commitment=False, + key_encipherment=True, + data_encipherment=False, + key_agreement=False, + key_cert_sign=True, # This is a CA + crl_sign=True, + encipher_only=False, + decipher_only=False + ), + critical=True, + ) + + logger.debug("Signing CA certificate") + ca_cert = builder.sign( + private_key=ca_private_key, + algorithm=hashes.SHA256(), + ) + + # Save CA certificate and key + with open(os.path.join(Config.get_config().certs_dir, Config.get_config().ca_cert), "wb") as f: + logger.debug(f"Saving CA certificate to {Config.get_config().ca_cert}") + f.write(ca_cert.public_bytes(serialization.Encoding.PEM)) + + with open(os.path.join(Config.get_config().certs_dir, Config.get_config().ca_key), "wb") as f: + logger.debug(f"Saving CA key to {Config.get_config().ca_key}") + f.write(ca_private_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption() + )) + + # CA generated, now generate server certificate + + ## Generate new certificate for domain + logger.debug("Generating private key for server") + server_key = rsa.generate_private_key( + public_exponent=65537, + key_size=2048, # 2048 bits is sufficient for domain certs + ) + + name = x509.Name([ + x509.NameAttribute(NameOID.COMMON_NAME, "CodeGate Server"), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "CodeGate"), + ]) + + # Add extended key usage extension + builder = x509.CertificateBuilder() + builder = builder.subject_name(name) + builder = builder.issuer_name(ca_cert.subject) + builder = builder.public_key(server_key.public_key()) + builder = builder.serial_number(x509.random_serial_number()) + builder = builder.not_valid_before(datetime.datetime.utcnow()) + builder = builder.not_valid_after( + datetime.datetime.utcnow() + datetime.timedelta(days=365) + ) + + # Add domain to SAN + builder = builder.add_extension( + x509.SubjectAlternativeName([x509.DNSName("localhost")]), + critical=False, + ) + + # Add extended key usage + builder = builder.add_extension( + x509.ExtendedKeyUsage([ + ExtendedKeyUsageOID.SERVER_AUTH, + ExtendedKeyUsageOID.CLIENT_AUTH, + ]), + critical=False, + ) + + # Basic constraints (not a CA) + builder = builder.add_extension( + x509.BasicConstraints(ca=False, path_length=None), + critical=True, + ) + + logger.debug("Signing server certificate") + server_cert = builder.sign( + private_key=ca_private_key, + algorithm=hashes.SHA256(), + ) + + # os.path.join(Config.get_config().server_key) + with open(os.path.join(Config.get_config().certs_dir, Config.get_config().server_cert), "wb") as f: + logger.debug(f"Saving server certificate to {Config.get_config().server_cert}") + f.write(server_cert.public_bytes(serialization.Encoding.PEM)) + + with open(os.path.join(Config.get_config().certs_dir, Config.get_config().server_key), "wb") as f: + logger.debug(f"Saving server key to {Config.get_config().server_key}") + f.write(server_key.private_bytes( + encoding=serialization.Encoding.PEM, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption() + )) + + # Print instructions for trusting the certificates + print("Certificates generated successfully in the 'certs' directory") + print("\nTo trust these certificates:") + print("\nOn macOS:") + print("`sudo security add-trusted-cert -d -r trustRoot -k /Library/Keychains/System.keychain certs/ca.crt") + print("\nOn Windows (PowerShell as Admin):") + print("Import-Certificate -FilePath \"certs\\ca.crt\" -CertStoreLocation Cert:\\LocalMachine\\Root") + print("\nOn Linux:") + print("sudo cp certs/ca.crt /usr/local/share/ca-certificates/codegate.crt") + print("sudo update-ca-certificates") + print("\nFor VSCode, add to settings.json:") + print('''{ + "http.proxy": "https://localhost:8990", + "http.proxySupport": "on", + "github.copilot.advanced": { + "debug.testOverrideProxyUrl": "https://localhost:8990", + "debug.overrideProxyUrl": "https://localhost:8990" + } +}''') + logger.debug("Certificates generated successfully") + return server_cert, server_key + + def create_ssl_context(self) -> ssl.SSLContext: + """Create SSL context with secure configuration""" + logger.debug("Creating SSL context fn: create_ssl_context") + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + logger.debug(f"Loading server certificate for ssl_context from: {Config.get_config().server_cert}") + ssl_context.load_cert_chain( + os.path.join(Config.get_config().certs_dir, Config.get_config().server_cert), + os.path.join(Config.get_config().certs_dir, Config.get_config().server_key)) + ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2 + ssl_context.options |= ( + ssl.OP_NO_SSLv2 | + ssl.OP_NO_SSLv3 | + ssl.OP_NO_COMPRESSION | + ssl.OP_CIPHER_SERVER_PREFERENCE + ) + ssl_context.set_ciphers('ECDHE+AESGCM:ECDHE+CHACHA20:DHE+AESGCM:DHE+CHACHA20') + logger.debug("SSL context created successfully") + return ssl_context + + def ensure_certificates_exist(self) -> None: + """Ensure SSL certificates exist, generate if they don't""" + logger.debug("Ensuring certificates exist. fn ensure_certificates_exist") + if not (os.path.exists(os.path.join(Config.get_config().certs_dir, Config.get_config().server_cert)) and os.path.exists(os.path.join(Config.get_config().certs_dir, Config.get_config().server_key))): + logger.debug("Certificates not found, generating new certificates") + self.generate_certificates() + else: + logger.debug(f"Certificates found at: {Config.get_config().server_cert} and {Config.get_config().server_key}") + + def get_ssl_context(self) -> ssl.SSLContext: + """Get SSL context with certificates""" + logger.debug("Getting SSL context fn: get_ssl_context") + self.ensure_certificates_exist() + return self.create_ssl_context() + + def get_cert_files(self) -> Tuple[str, str]: + """Get certificate and key file paths""" + logger.debug("Getting certificate and key file paths fn: get_cert_files") + return Config.get_config().server_cert, Config.get_config().server_key diff --git a/src/codegate/cli.py b/src/codegate/cli.py index f514e092..43ab3205 100644 --- a/src/codegate/cli.py +++ b/src/codegate/cli.py @@ -1,8 +1,12 @@ """Command-line interface for codegate.""" import sys +import asyncio +import signal from pathlib import Path from typing import Dict, Optional +from uvicorn.config import Config as UvicornConfig +from uvicorn.server import Server import click import structlog @@ -12,9 +16,68 @@ from codegate.db.connection import init_db_sync from codegate.server import init_app from codegate.storage.utils import restore_storage_backup +from codegate.providers.copilot.provider import CopilotProvider +from codegate.ca.codegate_ca import CertificateAuthority +logger = structlog.get_logger("codegate") + +class UvicornServer: + def __init__(self, config: UvicornConfig, server: Server): + self.server = server + self.config = config + self.port = config.port + self.host = config.host + self.log_level = config.log_level + self.log_config = None + self._startup_complete = asyncio.Event() + self._shutdown_event = asyncio.Event() + self._should_exit = False + + + + async def serve(self) -> None: + """Start the uvicorn server and handle shutdown gracefully.""" + logger.debug(f"Starting server on {self.host}:{self.port}") + + self.server = Server(config=self.config) + self.server.force_exit = True + + try: + self._startup_complete.set() + await self.server.serve() + except asyncio.CancelledError: + logger.info("Server received cancellation") + except Exception as e: + logger.exception("Unexpected error occurred during server execution", exc_info=e) + finally: + await self.cleanup() + + async def wait_startup_complete(self) -> None: + """Wait for the server to complete startup.""" + logger.debug("Waiting for server startup to complete") + await self._startup_complete.wait() + + async def cleanup(self) -> None: + """Cleanup server resources and ensure graceful shutdown.""" + logger.debug("Cleaning up server resources") + if not self._should_exit: + self._should_exit = True + logger.debug("Initiating server shutdown") + self._shutdown_event.set() + + if hasattr(self.server, 'shutdown'): + logger.debug("Shutting down server") + await self.server.shutdown() + + # Ensure all connections are closed + tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] + [task.cancel() for task in tasks] + + await asyncio.gather(*tasks, return_exceptions=True) + logger.debug("Server shutdown complete") def validate_port(ctx: click.Context, param: click.Parameter, value: int) -> int: + logger.debug(f"Validating port number: {value}") """Validate the port number is in valid range.""" if value is not None and not (1 <= value <= 65535): raise click.BadParameter("Port must be between 1 and 65535") @@ -61,6 +124,13 @@ def show_prompts(prompts: Optional[Path]) -> None: callback=validate_port, help="Port to listen on (default: 8989)", ) +@click.option( + "--proxy-port", + type=int, + default=None, + callback=validate_port, + help="Proxy port to listen on (default: 8990)", +) @click.option( "--host", type=str, @@ -127,6 +197,35 @@ def show_prompts(prompts: Optional[Path]) -> None: default="all-minilm-L6-v2-q5_k_m.gguf", help="Name of the model to use for embeddings", ) +@click.option( + "--certs-dir", + type=str, + default=None, + help="Directory for certificate files (default: ./certs)", +) +@click.option( + "--ca-cert", + type=str, + default=None, + help="CA certificate file name (default: ca.crt)", +) +@click.option( + "--ca-key", + type=str, + default=None, + help="CA key file name (default: ca.key)", +) +@click.option( + "--server-cert", + type=str, + default=None, + help="Server certificate file name (default: server.crt)", +) +@click.option( + "--server-key", + type=str, + default=None, + help="Server key file name (default: server.key)", @click.option( "--db-path", type=str, @@ -135,6 +234,7 @@ def show_prompts(prompts: Optional[Path]) -> None: ) def serve( port: Optional[int], + proxy_port: Optional[int], host: Optional[str], log_level: Optional[str], log_format: Optional[str], @@ -146,10 +246,14 @@ def serve( ollama_url: Optional[str], model_base_path: Optional[str], embedding_model: Optional[str], + certs_dir: Optional[str], + ca_cert: Optional[str], + ca_key: Optional[str], + server_cert: Optional[str], + server_key: Optional[str], db_path: Optional[str], ) -> None: """Start the codegate server.""" - logger = None try: # Create provider URLs dict from CLI options cli_provider_urls: Dict[str, str] = {} @@ -167,15 +271,56 @@ def serve( config_path=config, prompts_path=prompts, cli_port=port, + cli_proxy_port=proxy_port, cli_host=host, cli_log_level=log_level, cli_log_format=log_format, cli_provider_urls=cli_provider_urls, model_base_path=model_base_path, embedding_model=embedding_model, + certs_dir=certs_dir, + ca_cert=ca_cert, + ca_key=ca_key, + server_cert=server_cert, + server_key=server_key, db_path=db_path, ) + init_db_sync() + + + # Check certificates and create CA if necessary + logger.info("Checking certificates and creating CA our created") + ca = CertificateAuthority.get_instance() + ca.ensure_certificates_exist() + + # Set up event loop + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + app = init_app() + + # Run the server + try: + loop.run_until_complete(run_servers(cfg, app)) + except KeyboardInterrupt: + logger.info("Received keyboard interrupt, shutting down...") + finally: + loop.close() + + except ConfigurationError as e: + click.echo(f"Configuration error: {e}", err=True) + sys.exit(1) + except Exception as e: + if logger: + logger.exception("Unexpected error occurred") + click.echo(f"Error: {e}", err=True) + sys.exit(1) + + +async def run_servers(cfg: Config, app) -> None: + """Run the codegate server.""" + try: setup_logging(cfg.log_level, cfg.log_format) logger = structlog.get_logger("codegate") logger.info( @@ -183,12 +328,14 @@ def serve( extra={ "host": cfg.host, "port": cfg.port, + "proxy_port": cfg.proxy_port, "log_level": cfg.log_level.value, "log_format": cfg.log_format.value, "prompts_loaded": len(cfg.prompts.prompts), "provider_urls": cfg.provider_urls, "model_base_path": cfg.model_base_path, "embedding_model": cfg.embedding_model, + "certs_dir": cfg.certs_dir, "db_path": cfg.db_path, }, ) @@ -198,7 +345,8 @@ def serve( import uvicorn - uvicorn.run( + # Create Uvicorn configuration + uvicorn_config = UvicornConfig( app, host=cfg.host, port=cfg.port, @@ -206,19 +354,51 @@ def serve( log_config=None, # Default logging configuration ) - except KeyboardInterrupt: - if logger: - logger.info("Shutting down server") - except ConfigurationError as e: - click.echo(f"Configuration error: {e}", err=True) - sys.exit(1) - except Exception as e: - if logger: - logger.exception("Unexpected error occurred") - click.echo(f"Error: {e}", err=True) - sys.exit(1) + server = UvicornServer(uvicorn_config, Server(config=uvicorn_config)) + + # Initialize CopilotProvider and call run_proxy_server + copilot_provider = CopilotProvider(cfg) + # copilot_provider.run_proxy_server(cfg.proxy_port) + tasks = [ + asyncio.create_task(server.serve()), # Uvicorn server + asyncio.create_task(copilot_provider.run_proxy_server()) # Proxy server + ] + + # Create and start Uvicorn server + # server = UvicornServer(uvicorn_config, Server(config=uvicorn_config)) + + try: + await asyncio.gather(*tasks) + except asyncio.CancelledError: + logger.info("Server received cancellation") + except Exception as e: + logger.exception("Unexpected error occurred during server execution", exc_info=e) + finally: + await server.cleanup() + # Cleanup + for task in tasks: + if not task.done(): + task.cancel() + try: + await asyncio.gather(*tasks, return_exceptions=True) + except asyncio.CancelledError: + pass + + + # Set up signal handlers + # loop = asyncio.get_running_loop() + # for sig in (signal.SIGTERM, signal.SIGINT): + # loop.add_signal_handler(sig, lambda: asyncio.create_task(server.cleanup())) + + # # Start server + # await server.serve() + + except Exception as e: + logger.exception("Error running servers") + raise e + @cli.command() @click.option( "--backup-path", diff --git a/src/codegate/codegate_logging.py b/src/codegate/codegate_logging.py index 9a27a4f8..3d93cda1 100644 --- a/src/codegate/codegate_logging.py +++ b/src/codegate/codegate_logging.py @@ -1,7 +1,10 @@ import logging import sys +import json from enum import Enum from typing import Optional +from datetime import datetime +from typing import Any, Dict import structlog @@ -146,3 +149,53 @@ def setup_logging( "handlers": ["stdout", "stderr"], }, ) + + +def serialize_for_logging(obj: Any) -> Any: + """Serialize objects for logging, handling non-JSON serializable types""" + if hasattr(obj, '__dict__'): + return str(obj) + elif isinstance(obj, (datetime, bytes)): + return str(obj) + elif isinstance(obj, dict): + return {k: serialize_for_logging(v) for k, v in obj.items()} + elif isinstance(obj, (list, tuple)): + return [serialize_for_logging(item) for item in obj] + return obj + +def log_request(method: str, path: str, status_code: int, client: Any) -> None: + """Log HTTP request details""" + logger = logging.getLogger("proxy_pilot") + log_data = { + "timestamp": datetime.now().isoformat(), + "type": "request", + "method": method, + "path": path, + "status_code": status_code, + "client": serialize_for_logging(client) + } + logger.info(f"Request: {json.dumps(log_data, indent=2)}") + +def log_proxy_forward(target_url: str, method: str, status_code: int) -> None: + """Log proxy forwarding details""" + logger = logging.getLogger("proxy_pilot") + log_data = { + "timestamp": datetime.now().isoformat(), + "type": "proxy_forward", + "target_url": target_url, + "method": method, + "status_code": status_code + } + logger.info(f"Proxy Forward: {json.dumps(log_data, indent=2)}") + +def log_error(error_type: str, message: str, details: Dict = None) -> None: + """Log error details""" + logger = logging.getLogger("proxy_pilot") + log_data = { + "timestamp": datetime.now().isoformat(), + "type": "error", + "error_type": error_type, + "message": message, + "details": serialize_for_logging(details or {}) + } + logger.error(f"Error: {json.dumps(log_data, indent=2)}") diff --git a/src/codegate/config.py b/src/codegate/config.py index 6ac561bb..d8ad43fd 100644 --- a/src/codegate/config.py +++ b/src/codegate/config.py @@ -32,6 +32,7 @@ class Config: __config = None port: int = 8989 + proxy_port: int = 8990 # New proxy port configuration host: str = "localhost" log_level: LogLevel = LogLevel.INFO log_format: LogFormat = LogFormat.JSON @@ -43,6 +44,13 @@ class Config: embedding_model: str = "all-minilm-L6-v2-q5_k_m.gguf" db_path: Optional[str] = None + # Certificate configuration + certs_dir: str = "./certs" + ca_cert: str = "ca.crt" + ca_key: str = "ca.key" + server_cert: str = "server.crt" + server_key: str = "server.key" + # Provider URLs with defaults provider_urls: Dict[str, str] = field(default_factory=lambda: DEFAULT_PROVIDER_URLS.copy()) @@ -50,6 +58,9 @@ def __post_init__(self) -> None: """Validate configuration after initialization.""" if not isinstance(self.port, int) or not (1 <= self.port <= 65535): raise ConfigurationError("Port must be between 1 and 65535") + + if not isinstance(self.proxy_port, int) or not (1 <= self.proxy_port <= 65535): + raise ConfigurationError("Proxy port must be between 1 and 65535") if not isinstance(self.log_level, LogLevel): try: @@ -116,6 +127,7 @@ def from_file(cls, config_path: Union[str, Path]) -> "Config": return cls( port=config_data.get("port", cls.port), + proxy_port=config_data.get("proxy_port", cls.proxy_port), host=config_data.get("host", cls.host), log_level=config_data.get("log_level", cls.log_level.value), log_format=config_data.get("log_format", cls.log_format.value), @@ -125,6 +137,11 @@ def from_file(cls, config_path: Union[str, Path]) -> "Config": "chat_model_n_gpu_layers", cls.chat_model_n_gpu_layers ), embedding_model=config_data.get("embedding_model", cls.embedding_model), + certs_dir=config_data.get("certs_dir", cls.certs_dir), + ca_cert=config_data.get("ca_cert", cls.ca_cert), + ca_key=config_data.get("ca_key", cls.ca_key), + server_cert=config_data.get("server_cert", cls.server_cert), + server_key=config_data.get("server_key", cls.server_key), prompts=prompts_config, provider_urls=provider_urls, ) @@ -146,6 +163,8 @@ def from_env(cls) -> "Config": if "CODEGATE_APP_PORT" in os.environ: config.port = int(os.environ["CODEGATE_APP_PORT"]) + if "CODEGATE_APP_PROXY_PORT" in os.environ: + config.proxy_port = int(os.environ["CODEGATE_APP_PROXY_PORT"]) if "CODEGATE_APP_HOST" in os.environ: config.host = os.environ["CODEGATE_APP_HOST"] if "CODEGATE_APP_LOG_LEVEL" in os.environ: @@ -157,6 +176,18 @@ def from_env(cls) -> "Config": os.environ["CODEGATE_PROMPTS_FILE"] ) # noqa: E501 + # Load certificate configuration from environment + if "CODEGATE_CERTS_DIR" in os.environ: + config.certs_dir = os.environ["CODEGATE_CERTS_DIR"] + if "CODEGATE_CA_CERT" in os.environ: + config.ca_cert = os.environ["CODEGATE_CA_CERT"] + if "CODEGATE_CA_KEY" in os.environ: + config.ca_key = os.environ["CODEGATE_CA_KEY"] + if "CODEGATE_SERVER_CERT" in os.environ: + config.server_cert = os.environ["CODEGATE_SERVER_CERT"] + if "CODEGATE_SERVER_KEY" in os.environ: + config.server_key = os.environ["CODEGATE_SERVER_KEY"] + # Load provider URLs from environment variables for provider in DEFAULT_PROVIDER_URLS.keys(): env_var = f"CODEGATE_PROVIDER_{provider.upper()}_URL" @@ -173,12 +204,18 @@ def load( config_path: Optional[Union[str, Path]] = None, prompts_path: Optional[Union[str, Path]] = None, cli_port: Optional[int] = None, + cli_proxy_port: Optional[int] = None, cli_host: Optional[str] = None, cli_log_level: Optional[str] = None, cli_log_format: Optional[str] = None, cli_provider_urls: Optional[Dict[str, str]] = None, model_base_path: Optional[str] = None, embedding_model: Optional[str] = None, + certs_dir: Optional[str] = None, + ca_cert: Optional[str] = None, + ca_key: Optional[str] = None, + server_cert: Optional[str] = None, + server_key: Optional[str] = None, db_path: Optional[str] = None, ) -> "Config": """Load configuration with priority resolution. @@ -193,12 +230,18 @@ def load( config_path: Optional path to config file prompts_path: Optional path to prompts file cli_port: Optional CLI port override + cli_proxy_port: Optional CLI proxy port override cli_host: Optional CLI host override cli_log_level: Optional CLI log level override cli_log_format: Optional CLI log format override cli_provider_urls: Optional dict of provider URLs from CLI model_base_path: Optional path to model base directory embedding_model: Optional name of the model to use for embeddings + certs_dir: Optional path to certificates directory + ca_cert: Optional path to CA certificate + ca_key: Optional path to CA key + server_cert: Optional path to server certificate + server_key: Optional path to server key db_path: Optional path to the SQLite database file Returns: @@ -222,6 +265,8 @@ def load( env_config = cls.from_env() if "CODEGATE_APP_PORT" in os.environ: config.port = env_config.port + if "CODEGATE_APP_PROXY_PORT" in os.environ: + config.proxy_port = env_config.proxy_port if "CODEGATE_APP_HOST" in os.environ: config.host = env_config.host if "CODEGATE_APP_LOG_LEVEL" in os.environ: @@ -234,6 +279,16 @@ def load( config.model_base_path = env_config.model_base_path if "CODEGATE_EMBEDDING_MODEL" in os.environ: config.embedding_model = env_config.embedding_model + if "CODEGATE_CERTS_DIR" in os.environ: + config.certs_dir = env_config.certs_dir + if "CODEGATE_CA_CERT" in os.environ: + config.ca_cert = env_config.ca_cert + if "CODEGATE_CA_KEY" in os.environ: + config.ca_key = env_config.ca_key + if "CODEGATE_SERVER_CERT" in os.environ: + config.server_cert = env_config.server_cert + if "CODEGATE_SERVER_KEY" in os.environ: + config.server_key = env_config.server_key # Override provider URLs from environment for provider, url in env_config.provider_urls.items(): @@ -242,6 +297,8 @@ def load( # Override with CLI arguments if cli_port is not None: config.port = cli_port + if cli_proxy_port is not None: + config.proxy_port = cli_proxy_port if cli_host is not None: config.host = cli_host if cli_log_level is not None: @@ -256,6 +313,16 @@ def load( config.model_base_path = model_base_path if embedding_model is not None: config.embedding_model = embedding_model + if certs_dir is not None: + config.certs_dir = certs_dir + if ca_cert is not None: + config.ca_cert = ca_cert + if ca_key is not None: + config.ca_key = ca_key + if server_cert is not None: + config.server_cert = server_cert + if server_key is not None: + config.server_key = server_key if db_path is not None: config.db_path = db_path diff --git a/src/codegate/providers/copilot/__init__.py b/src/codegate/providers/copilot/__init__.py new file mode 100644 index 00000000..d342fb0e --- /dev/null +++ b/src/codegate/providers/copilot/__init__.py @@ -0,0 +1,5 @@ +from codegate.providers.copilot.provider import CopilotProvider + +__all__ = [ + "CopilotProvider", +] diff --git a/src/codegate/providers/copilot/logging.py b/src/codegate/providers/copilot/logging.py new file mode 100644 index 00000000..264fcd22 --- /dev/null +++ b/src/codegate/providers/copilot/logging.py @@ -0,0 +1,79 @@ +import logging +import json +from datetime import datetime +from typing import Any, Dict + +def setup_logging() -> logging.Logger: + """Setup logging configuration""" + + # Create logger + logger = logging.getLogger("proxy_pilot") + logger.setLevel(logging.DEBUG) + + # Create formatter + formatter = logging.Formatter( + '%(asctime)s - %(name)s - %(levelname)s - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S%f' + ) + + # Create console handler + console_handler = logging.StreamHandler() + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + return logger + +def serialize_for_logging(obj: Any) -> Any: + """Serialize objects for logging, handling non-JSON serializable types""" + if hasattr(obj, '__dict__'): + return str(obj) + elif isinstance(obj, (datetime, bytes)): + return str(obj) + elif isinstance(obj, dict): + return {k: serialize_for_logging(v) for k, v in obj.items()} + elif isinstance(obj, (list, tuple)): + return [serialize_for_logging(item) for item in obj] + return obj + +def log_request(method: str, path: str, status_code: int, client: Any) -> None: + """Log HTTP request details""" + logger = logging.getLogger("proxy_pilot") + log_data = { + "timestamp": datetime.now().isoformat(), + "type": "request", + "method": method, + "path": path, + "status_code": status_code, + "client": serialize_for_logging(client) + } + logger.info(f"Request: {json.dumps(log_data, indent=2)}") + +def log_proxy_forward(target_url: str, method: str, status_code: int) -> None: + """Log proxy forwarding details""" + logger = logging.getLogger("proxy_pilot") + log_data = { + "timestamp": datetime.now().isoformat(), + "type": "proxy_forward", + "target_url": target_url, + "method": method, + "status_code": status_code + } + logger.info(f"Proxy Forward: {json.dumps(log_data, indent=2)}") + +def log_error(error_type: str, message: str, details: Dict = None) -> None: + """Log error details""" + logger = logging.getLogger("proxy_pilot") + log_data = { + "timestamp": datetime.now().isoformat(), + "type": "error", + "error_type": error_type, + "message": message, + "details": serialize_for_logging(details or {}) + } + logger.error(f"Error: {json.dumps(log_data, indent=2)}") + +# Initialize logger +logger = setup_logging() + + +logger \ No newline at end of file diff --git a/src/codegate/providers/copilot/mapping.py b/src/codegate/providers/copilot/mapping.py new file mode 100644 index 00000000..9bda5139 --- /dev/null +++ b/src/codegate/providers/copilot/mapping.py @@ -0,0 +1,42 @@ +from typing import List, Dict +from pydantic import BaseModel, HttpUrl +from pydantic_settings import BaseSettings + +class CopilotProxyRoute(BaseModel): + """Pydantic model for proxy route validation""" + path: str + target: HttpUrl + +class CoPilotMappings(BaseSettings): + + # Proxy routes configuration + # This code is to ensure that incoming proxy requests are routed to the correct target + PROXY_ROUTES: List[tuple[str, str]] = [ + ("/github/login", "https://github.com/login"), + ("/api/github/user", "https://api.github.com"), + ("/api/github/copilot", "https://api.github.com/copilot_internal"), + ("/copilot/telemetry", "https://copilot-telemetry.githubusercontent.com"), + ("/vscode/ab", "https://default.exp-tas.com"), + ("/copilot/proxy", "https://copilot-proxy.githubusercontent.com"), + ("/origin-tracker", "https://origin-tracker.githubusercontent.com"), + ("/copilot/suggestions", "https://githubcopilot.com"), + ("/copilot_internal/user", "https://api.github.com"), + ("/copilot_internal/v2/token", "https://api.github.com"), + ("/telemety", "https://copilot-telemetry.githubusercontent.com"), + ("/", "https://github.com"), + ("/login/oauth/access_token", "https://github.com/login/oauth/access_token"), + ("/api/copilot", "https://api.github.com/copilot_internal"), + ("/api/copilot_internal", "https://api.github.com/copilot_internal"), + ("/v1/completions", "https://copilot-proxy.githubusercontent.com/v1/completions"), + ("/v1", "https://copilot-proxy.githubusercontent.com/v1"), + ] + +# Create settings instance +mappings = CoPilotMappings() + +# Convert routes to validated ProxyRoute objects +VALIDATED_ROUTES: List[CopilotProxyRoute] = [ + CopilotProxyRoute(path=path, target=target) + for path, target in mappings.PROXY_ROUTES +] + diff --git a/src/codegate/providers/copilot/provider.py b/src/codegate/providers/copilot/provider.py new file mode 100644 index 00000000..ebae6947 --- /dev/null +++ b/src/codegate/providers/copilot/provider.py @@ -0,0 +1,486 @@ +import asyncio +import re +import ssl +from typing import Dict, Optional, Tuple, Union +from urllib.parse import unquote, urljoin, urlparse + +import structlog + +from codegate.config import Config +from codegate.ca.codegate_ca import CertificateAuthority +from codegate.providers.copilot.mapping import VALIDATED_ROUTES + +logger = structlog.get_logger("codegate") + +# Increase buffer sizes +MAX_BUFFER_SIZE = 10 * 1024 * 1024 # 10MB +CHUNK_SIZE = 64 * 1024 # 64KB + +class CopilotProvider(asyncio.Protocol): + def __init__(self, loop): + logger.debug("Initializing CopilotProvider class: CopilotProvider") + self.loop = loop + self.transport: Optional[asyncio.Transport] = None + self.target_transport: Optional[asyncio.Transport] = None + self.peername: Optional[Tuple[str, int]] = None + self.buffer = bytearray() + self.target_host: Optional[str] = None + self.target_port: Optional[int] = None + self.handshake_done = False + self.is_connect = False + self.content_length = 0 + self.headers_parsed = False + self.method = None + self.path = None + self.version = None + self.headers = [] + self.target = None + self.original_path = None + self.ssl_context = None + self.proxy_ep = None + self.decrypted_data = bytearray() + # Get the singleton instance of CertificateAuthority + self.ca = CertificateAuthority.get_instance() + + def connection_made(self, transport: asyncio.Transport): + logger.debug("Client connected fn: connection_made") + self.transport = transport + self.peername = transport.get_extra_info('peername') + logger.debug(f"Client connected from {self.peername}") + + def extract_path(self, full_path: str) -> str: + logger.debug(f"Extracting path from {full_path} fn: extract_path") + if full_path.startswith('http://') or full_path.startswith('https://'): + parsed = urlparse(full_path) + path = parsed.path + if parsed.query: + path = f"{path}?{parsed.query}" + return path.lstrip('/') + elif full_path.startswith('/'): + return full_path.lstrip('/') + return full_path + + def get_headers(self) -> Dict[str, str]: + """Get request headers as a dictionary""" + logger.debug("Getting headers as dictionary fn: get_headers") + headers_dict = {} + + try: + if b'\r\n\r\n' not in self.buffer: + return {} + + headers_end = self.buffer.index(b'\r\n\r\n') + headers = self.buffer[:headers_end].split(b'\r\n')[1:] # Skip request line + + for header in headers: + try: + name, value = header.decode('utf-8').split(':', 1) + headers_dict[name.strip().lower()] = value.strip() + except ValueError: + continue + + return headers_dict + except Exception as e: + logger.error(f"Error getting headers: {e}") + return {} + + def parse_headers(self) -> bool: + logger.debug("Parsing headers fn: parse_headers") + try: + if b'\r\n\r\n' not in self.buffer: + return False + + headers_end = self.buffer.index(b'\r\n\r\n') + headers = self.buffer[:headers_end].split(b'\r\n') + + request = headers[0].decode('utf-8') + self.method, full_path, self.version = request.split(' ') + + self.original_path = full_path + + if self.method == 'CONNECT': + logger.debug(f"CONNECT request to {full_path}") + self.target = full_path + self.path = "" + else: + logger.debug(f"Request: {self.method} {full_path} {self.version}") + self.path = self.extract_path(full_path) + + self.headers = [header.decode('utf-8') for header in headers[1:]] + return True + except Exception as e: + logger.error(f"Error parsing headers: {e}") + return False + + def log_decrypted_data(self, data: bytes, direction: str): + ''' + Uncomment to log data from payload + ''' + try: + # decoded = data.decode('utf-8') + # logger.debug(f"=== Decrypted {direction} Data ===") + # logger.debug(decoded) + # logger.debug("=" * 40) + pass + except UnicodeDecodeError: + # pass + # logger.debug(f"=== Decrypted {direction} Data (hex) ===") + # logger.debug(data.hex()) + # logger.debug("=" * 40) + pass + + async def handle_http_request(self): + logger.debug("Handling HTTP request fn: handle_http_request") + logger.debug("=" * 40) + logger.debug(f"Method: {self.method}") + logger.debug(f"Searched Path: {self.path} in target URL") + try: + # Extract proxy endpoint from authorization header if present + headers_dict = self.get_headers() + auth_header = headers_dict.get('authorization', '') + if auth_header: + match = re.search(r"proxy-ep=([^;]+)", auth_header) + if match: + self.proxy_ep = match.group(1) + logger.debug(f"Extracted proxy-ep value: {self.proxy_ep}") + + # Check if the proxy_ep includes a scheme + parsed_proxy_ep = urlparse(self.proxy_ep) + if not parsed_proxy_ep.scheme: + # Default to https if no scheme is provided + self.proxy_ep = f"https://{self.proxy_ep}" + logger.debug(f"Added default scheme to proxy-ep: {self.proxy_ep}") + + target_url = f"{self.proxy_ep}/{self.path}" + else: + target_url = await self.get_target_url(self.path) + else: + target_url = await self.get_target_url(self.path) + + if not target_url: + self.send_error_response(404, b"Not Found") + return + logger.debug(f"Target URL: {target_url}") + + parsed_url = urlparse(target_url) + logger.debug(f"Parsed URL {parsed_url}") + + self.target_host = parsed_url.hostname + self.target_port = parsed_url.port or (443 if parsed_url.scheme == 'https' else 80) + + target_protocol = CopilotProxyTargetProtocol(self) + logger.debug(f"Connecting to {self.target_host}:{self.target_port}") + await self.loop.create_connection( + lambda: target_protocol, + self.target_host, + self.target_port, + ssl=parsed_url.scheme == 'https' + ) + + has_host = False + new_headers = [] + + for header in self.headers: + if header.lower().startswith('host:'): + has_host = True + new_headers.append(f"Host: {self.target_host}") + else: + new_headers.append(header) + + if not has_host: + new_headers.append(f"Host: {self.target_host}") + + request_line = f"{self.method} /{self.path} {self.version}\r\n".encode() + logger.debug(f"Request Line: {request_line}") + header_block = '\r\n'.join(new_headers).encode() + headers = request_line + header_block + b'\r\n\r\n' + + if self.target_transport: + logger.debug("=" * 40) + self.log_decrypted_data(headers, "Request") + self.target_transport.write(headers) + logger.debug("=" * 40) + + body_start = self.buffer.index(b'\r\n\r\n') + 4 + body = self.buffer[body_start:] + + if body: + self.log_decrypted_data(body, "Request Body") + + for i in range(0, len(body), CHUNK_SIZE): + chunk = body[i:i + CHUNK_SIZE] + self.target_transport.write(chunk) + else: + logger.debug("=" * 40) + logger.error("Target transport not available") + logger.debug("=" * 40) + self.send_error_response(502, b"Failed to establish target connection") + + except Exception as e: + logger.error(f"Error handling HTTP request: {e}") + self.send_error_response(502, str(e).encode()) + + def _check_buffer_size(self, new_data: bytes) -> bool: + """Check if adding new data would exceed the maximum buffer size""" + return len(self.buffer) + len(new_data) <= MAX_BUFFER_SIZE + + def _handle_parsed_headers(self) -> None: + """Handle the request based on parsed headers""" + if self.method == 'CONNECT': + logger.debug("Handling CONNECT request") + self.handle_connect() + else: + logger.debug("Handling HTTP request") + asyncio.create_task(self.handle_http_request()) + + def _forward_data_to_target(self, data: bytes) -> None: + """Forward data to target if connection is established""" + if self.target_transport and not self.target_transport.is_closing(): + self.log_decrypted_data(data, "Client to Server") + self.target_transport.write(data) + + def data_received(self, data: bytes) -> None: + """Handle received data from the client""" + logger.debug(f"Data received from {self.peername} fn: data_received") + + try: + # Check buffer size limit + if not self._check_buffer_size(data): + logger.error("Request exceeds maximum buffer size") + self.send_error_response(413, b"Request body too large") + return + + # Append new data to buffer + self.buffer.extend(data) + + if not self.headers_parsed: + # Try to parse headers + self.headers_parsed = self.parse_headers() + if not self.headers_parsed: + return + + # Handle the request based on parsed headers + self._handle_parsed_headers() + else: + # Forward data to target if headers are already parsed + self._forward_data_to_target(data) + + except asyncio.CancelledError: + logger.warning("Operation cancelled") + raise + except Exception as e: + logger.error(f"Error processing received data: {e}") + self.send_error_response(502, str(e).encode()) + + def handle_connect(self): + ''' + This where requests are sent directly via the tunnel created during + a CONNECT request. This is where the SSL context is created and the + internal connection is made to the target host. + + We do not need to do a URL to mapping, as this passes through the + tunnel with a FQDN already set by the source (client) request. + ''' + try: + path = unquote(self.target) + if ':' in path: + self.target_host, port = path.split(':') + self.target_port = int(port) + logger.debug("=" * 40) + logger.debug(f"CONNECT request to {self.target_host}:{self.target_port}") + logger.debug("Headers:") + for header in self.headers: + logger.debug(f" {header}") + + logger.debug("=" * 40) + cert_path, key_path = self.ca.get_domain_certificate(self.target_host) + + logger.debug(f"Setting up SSL context for {self.target_host}") + logger.debug(f"Using certificate: {cert_path}") + logger.debug(f"Using key: {key_path}") + + self.ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + self.ssl_context.load_cert_chain(cert_path, key_path) + self.ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2 + + self.is_connect = True + logger.debug("CONNECT handshake complete") + asyncio.create_task(self.connect_to_target()) + self.handshake_done = True + else: + logger.error(f"Invalid CONNECT path: {path}") + self.send_error_response(400, b"Invalid CONNECT path") + except Exception as e: + logger.error(f"Error handling CONNECT: {e}") + self.send_error_response(502, str(e).encode()) + + def send_error_response(self, status: int, message: bytes): + logger.debug(f"Sending error response: {status} {message} fn: send_error_response") + response = ( + f"HTTP/1.1 {status} {self.get_status_text(status)}\r\n" + f"Content-Length: {len(message)}\r\n" + f"Content-Type: text/plain\r\n" + f"\r\n" + ).encode() + message + if self.transport and not self.transport.is_closing(): + self.transport.write(response) + self.transport.close() + + def get_status_text(self, status: int) -> str: + logger.debug(f"Getting status text for {status} fn: get_status_text") + status_texts = { + 400: "Bad Request", + 404: "Not Found", + 413: "Request Entity Too Large", + 502: "Bad Gateway" + } + return status_texts.get(status, "Error") + + async def connect_to_target(self): + logger.debug(f"Connecting to target {self.target_host}:{self.target_port} fn: connect_to_target") + try: + if not self.target_host or not self.target_port: + raise ValueError("Target host and port not set") + + logger.debug(f"Attempting to connect to {self.target_host}:{self.target_port}") + + # Create SSL context for target connection + logger.debug("Creating SSL context for target connection") + target_ssl_context = ssl.create_default_context() + # Don't verify certificates when connecting to target + target_ssl_context.check_hostname = False + target_ssl_context.verify_mode = ssl.CERT_NONE + + # Connect directly to target host + logger.debug(f"Connecting to {self.target_host}:{self.target_port}") + target_protocol = CopilotProxyTargetProtocol(self) + transport, _ = await self.loop.create_connection( + lambda: target_protocol, + self.target_host, + self.target_port, + ssl=target_ssl_context + ) + + logger.debug(f"Successfully connected to {self.target_host}:{self.target_port}") + + # Send 200 Connection Established + if self.transport and not self.transport.is_closing(): + logger.debug("Sending 200 Connection Established response") + self.transport.write( + b'HTTP/1.1 200 Connection Established\r\n' + b'Proxy-Agent: ProxyPilot\r\n' + b'Connection: keep-alive\r\n\r\n' + ) + + # Upgrade client connection to SSL + logger.debug("Upgrading client connection to SSL") + transport = await self.loop.start_tls( + self.transport, + self, + self.ssl_context, + server_side=True + ) + self.transport = transport + + except Exception as e: + logger.error(f"Failed to connect to target {self.target_host}:{self.target_port}: {e}") + self.send_error_response(502, str(e).encode()) + + def connection_lost(self, exc): + logger.debug(f"Connection lost from {self.peername} fn: connection_lost") + logger.debug(f"Client disconnected from {self.peername}") + if self.target_transport and not self.target_transport.is_closing(): + self.target_transport.close() + + @classmethod + async def create_proxy_server(cls, host: str, port: int, ssl_context: Optional[ssl.SSLContext] = None): + logger.debug(f"Creating proxy server on {host}:{port} fn: create_proxy_server") + loop = asyncio.get_event_loop() + + def create_protocol(): + logger.debug("Creating protocol for proxy server fn: create_protocol") + return cls(loop) + + logger.debug(f"Starting proxy server on {host}:{port}") + server = await loop.create_server( + create_protocol, + host, + port, + ssl=ssl_context, + reuse_port=True, + start_serving=True + ) + + logger.debug(f"Proxy server running on {host}:{port}") + return server + + @classmethod + async def run_proxy_server(cls): + logger.debug("Running proxy server fn: run_proxy_server") + try: + # Get the singleton instance of CertificateAuthority + ca = CertificateAuthority.get_instance() + logger.debug("Creating SSL context for proxy server") + ssl_context = ca.create_ssl_context() + server = await cls.create_proxy_server( + Config.get_config().host, + Config.get_config().proxy_port, + ssl_context + ) + logger.debug("Proxy server created") + async with server: + await server.serve_forever() + + except Exception as e: + logger.error(f"Proxy server error: {e}") + raise + + @classmethod + async def get_target_url(cls, path: str) -> Optional[str]: + """Get target URL for the given path""" + logger.debug(f"Attempting to get target URL for path: {path} fn: get_target_url") + + logger.debug("=" * 40) + logger.debug("Validated routes:") + for route in VALIDATED_ROUTES: + if path == route.path: + logger.debug(f" {route.path} -> {route.target}") + logger.debug(f"Found exact path match: {path} -> {route.target}") + return str(route.target) + + # Then check for prefix match + for route in VALIDATED_ROUTES: + # For prefix matches, keep the rest of the path + remaining_path = path[len(route.path):] + logger.debug(f"Remaining path: {remaining_path}") + # Make sure we don't end up with double slashes + if remaining_path and remaining_path.startswith('/'): + remaining_path = remaining_path[1:] + target = urljoin(str(route.target), remaining_path) + logger.debug(f"Found prefix match: {path} -> {target} (using route {route.path} -> {route.target})") + return target + + logger.warning(f"No route found for path: {path}") + return None + +class CopilotProxyTargetProtocol(asyncio.Protocol): + def __init__(self, proxy: CopilotProvider): + logger.debug("Initializing CopilotProxyTargetProtocol class: CopilotProxyTargetProtocol") + self.proxy = proxy + self.transport: Optional[asyncio.Transport] = None + + def connection_made(self, transport: asyncio.Transport): + logger.debug(f"Connection made to target {self.proxy.target_host}:{self.proxy.target_port}") + self.transport = transport + self.proxy.target_transport = transport + + def data_received(self, data: bytes): + logger.debug(f"Data received from target {self.proxy.target_host}:{self.proxy.target_port}") + if self.proxy.transport and not self.proxy.transport.is_closing(): + self.proxy.log_decrypted_data(data, "Server to Client") + self.proxy.transport.write(data) + + def connection_lost(self, exc): + logger.debug(f"Connection lost from target {self.proxy.target_host}:{self.proxy.target_port}") + if self.proxy.transport and not self.proxy.transport.is_closing(): + self.proxy.transport.close() diff --git a/tests/test_cli.py b/tests/test_cli.py index 432721f4..69c39b2d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,13 +1,16 @@ """Tests for the CLI module.""" +import asyncio from pathlib import Path -from typing import Any -from unittest.mock import MagicMock, patch +from typing import Any, AsyncGenerator +import signal +from unittest.mock import MagicMock, patch, AsyncMock import pytest from click.testing import CliRunner +from uvicorn.config import Config as UvicornConfig -from codegate.cli import cli +from codegate.cli import cli, UvicornServer from codegate.codegate_logging import LogFormat, LogLevel from codegate.config import DEFAULT_PROVIDER_URLS @@ -44,22 +47,99 @@ def temp_config_file(tmp_path: Path) -> Path: host: localhost log_level: DEBUG log_format: JSON +certs_dir: "./test-certs" +ca_cert: "test-ca.crt" +ca_key: "test-ca.key" +server_cert: "test-server.crt" +server_key: "test-server.key" """ ) return config_file +@pytest.fixture +async def mock_uvicorn_server() -> AsyncGenerator[MagicMock, None]: + """Create a mock Uvicorn server.""" + config = UvicornConfig(app=MagicMock(), host="localhost", port=8989) + server = MagicMock() + server.serve = AsyncMock() + server.shutdown = AsyncMock() + yield UvicornServer(config=config, server=server) + + def test_cli_version(cli_runner: CliRunner) -> None: """Test CLI version command.""" result = cli_runner.invoke(cli, ["--version"]) assert result.exit_code == 0 +@pytest.mark.asyncio +async def test_uvicorn_server_serve(mock_uvicorn_server: UvicornServer) -> None: + """Test UvicornServer serve method.""" + # Start server in background task + server_task = asyncio.create_task(mock_uvicorn_server.serve()) + + # Wait for startup to complete + await mock_uvicorn_server.wait_startup_complete() + + # Verify server started + assert mock_uvicorn_server.server.serve.called + + # Cleanup + await mock_uvicorn_server.cleanup() + await server_task + + +@pytest.mark.asyncio +async def test_uvicorn_server_cleanup(mock_uvicorn_server: UvicornServer) -> None: + """Test UvicornServer cleanup method.""" + # Start server + server_task = asyncio.create_task(mock_uvicorn_server.serve()) + await mock_uvicorn_server.wait_startup_complete() + + # Trigger cleanup + await mock_uvicorn_server.cleanup() + + # Verify shutdown was called + assert mock_uvicorn_server.server.shutdown.called + assert mock_uvicorn_server._shutdown_event.is_set() + + await server_task + + +@pytest.mark.asyncio +async def test_uvicorn_server_signal_handling(mock_uvicorn_server: UvicornServer) -> None: + """Test signal handling in UvicornServer.""" + # Mock signal handlers + with patch('asyncio.get_running_loop') as mock_loop: + mock_loop_instance = MagicMock() + mock_loop.return_value = mock_loop_instance + + # Start server + server_task = asyncio.create_task(mock_uvicorn_server.serve()) + await mock_uvicorn_server.wait_startup_complete() + + # Simulate SIGTERM + mock_loop_instance.add_signal_handler.assert_any_call( + signal.SIGTERM, + pytest.approx(type(lambda: None)) # Check if a callable was passed + ) + + # Simulate SIGINT + mock_loop_instance.add_signal_handler.assert_any_call( + signal.SIGINT, + pytest.approx(type(lambda: None)) # Check if a callable was passed + ) + + await mock_uvicorn_server.cleanup() + await server_task + + def test_serve_default_options( cli_runner: CliRunner, mock_logging: Any, mock_setup_logging: Any ) -> None: """Test serve command with default options.""" - with patch("uvicorn.run") as mock_run: + with patch("codegate.cli.run_servers") as mock_run: logger_instance = MagicMock() mock_logging.return_value = logger_instance result = cli_runner.invoke(cli, ["serve"]) @@ -68,7 +148,7 @@ def test_serve_default_options( mock_setup_logging.assert_called_once_with(LogLevel.INFO, LogFormat.JSON) mock_logging.assert_called_with("codegate") - # validate only a subset of the expected extra arguments, as image provides more + # validate only a subset of the expected extra arguments expected_extra = { "host": "localhost", "port": 8989, @@ -76,6 +156,7 @@ def test_serve_default_options( "log_format": "JSON", "prompts_loaded": 7, "provider_urls": DEFAULT_PROVIDER_URLS, + "certs_dir": "./certs", # Default certificate directory } # Retrieve the actual call arguments @@ -93,7 +174,7 @@ def test_serve_custom_options( cli_runner: CliRunner, mock_logging: Any, mock_setup_logging: Any ) -> None: """Test serve command with custom options.""" - with patch("uvicorn.run") as mock_run: + with patch("codegate.cli.run_servers") as mock_run: logger_instance = MagicMock() mock_logging.return_value = logger_instance result = cli_runner.invoke( @@ -108,6 +189,16 @@ def test_serve_custom_options( "DEBUG", "--log-format", "TEXT", + "--certs-dir", + "./custom-certs", + "--ca-cert", + "custom-ca.crt", + "--ca-key", + "custom-ca.key", + "--server-cert", + "custom-server.crt", + "--server-key", + "custom-server.key", ], ) @@ -125,6 +216,7 @@ def test_serve_custom_options( "log_format": "TEXT", "prompts_loaded": 7, # Default prompts are loaded "provider_urls": DEFAULT_PROVIDER_URLS, + "certs_dir": "./custom-certs", } # Check if one of the calls matches the expected subset @@ -153,7 +245,7 @@ def test_serve_with_config_file( cli_runner: CliRunner, mock_logging: Any, temp_config_file: Path, mock_setup_logging: Any ) -> None: """Test serve command with config file.""" - with patch("uvicorn.run") as mock_run: + with patch("codegate.cli.run_servers") as mock_run: logger_instance = MagicMock() mock_logging.return_value = logger_instance result = cli_runner.invoke(cli, ["serve", "--config", str(temp_config_file)]) @@ -172,6 +264,7 @@ def test_serve_with_config_file( "log_format": "JSON", "prompts_loaded": 7, # Default prompts are loaded "provider_urls": DEFAULT_PROVIDER_URLS, + "certs_dir": "./test-certs", # From config file } # Check if one of the calls matches the expected subset @@ -197,7 +290,7 @@ def test_serve_priority_resolution( mock_setup_logging: Any, ) -> None: """Test serve command respects configuration priority.""" - with patch("uvicorn.run") as mock_run: + with patch("codegate.cli.run_servers") as mock_run: logger_instance = MagicMock() mock_logging.return_value = logger_instance result = cli_runner.invoke( @@ -214,6 +307,16 @@ def test_serve_priority_resolution( "ERROR", "--log-format", "TEXT", + "--certs-dir", + "./cli-certs", + "--ca-cert", + "cli-ca.crt", + "--ca-key", + "cli-ca.key", + "--server-cert", + "cli-server.crt", + "--server-key", + "cli-server.key", ], ) @@ -231,6 +334,56 @@ def test_serve_priority_resolution( "log_format": "TEXT", "prompts_loaded": 7, # Default prompts are loaded "provider_urls": DEFAULT_PROVIDER_URLS, + "certs_dir": "./cli-certs", # CLI args override config file + } + + # Check if one of the calls matches the expected subset + assert any( + all(expected_extra[k] == actual_extra.get(k) for k in expected_extra) + for actual_extra in calls + ) + mock_run.assert_called_once() + + +def test_serve_certificate_options( + cli_runner: CliRunner, mock_logging: Any, mock_setup_logging: Any +) -> None: + """Test serve command with certificate options.""" + with patch("codegate.cli.run_servers") as mock_run: + logger_instance = MagicMock() + mock_logging.return_value = logger_instance + result = cli_runner.invoke( + cli, + [ + "serve", + "--certs-dir", + "./custom-certs", + "--ca-cert", + "custom-ca.crt", + "--ca-key", + "custom-ca.key", + "--server-cert", + "custom-server.crt", + "--server-key", + "custom-server.key", + ], + ) + + assert result.exit_code == 0 + mock_setup_logging.assert_called_once_with(LogLevel.INFO, LogFormat.JSON) + mock_logging.assert_called_with("codegate") + + # Retrieve the actual call arguments + calls = [call[1]["extra"] for call in logger_instance.info.call_args_list] + + expected_extra = { + "host": "localhost", + "port": 8989, + "log_level": "INFO", + "log_format": "JSON", + "prompts_loaded": 6, + "provider_urls": DEFAULT_PROVIDER_URLS, + "certs_dir": "./custom-certs", } # Check if one of the calls matches the expected subset diff --git a/tests/test_config.py b/tests/test_config.py index 3f6ab585..1b9bf4b9 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -16,6 +16,12 @@ def test_default_config(default_config: Config) -> None: assert default_config.log_level == LogLevel.INFO assert default_config.log_format == LogFormat.JSON assert default_config.provider_urls == DEFAULT_PROVIDER_URLS + # Test default certificate configuration + assert default_config.certs_dir == "./certs" + assert default_config.ca_cert == "ca.crt" + assert default_config.ca_key == "ca.key" + assert default_config.server_cert == "server.crt" + assert default_config.server_key == "server.key" def test_config_from_file(temp_config_file: Path) -> None: @@ -64,12 +70,22 @@ def test_config_priority_resolution(temp_config_file: Path, env_vars: None) -> N cli_log_level="WARNING", cli_log_format="TEXT", cli_provider_urls={"vllm": "https://custom.vllm.server"}, + certs_dir="./custom-certs", + ca_cert="custom-ca.crt", + ca_key="custom-ca.key", + server_cert="custom-server.crt", + server_key="custom-server.key", ) assert config.port == 8080 assert config.host == "example.com" assert config.log_level == LogLevel.WARNING assert config.log_format == LogFormat.TEXT assert config.provider_urls["vllm"] == "https://custom.vllm.server" + assert config.certs_dir == "./custom-certs" + assert config.ca_cert == "custom-ca.crt" + assert config.ca_key == "custom-ca.key" + assert config.server_cert == "custom-server.crt" + assert config.server_key == "custom-server.key" # Env vars should override config file config = Config.load(config_path=temp_config_file) @@ -116,6 +132,54 @@ def test_provider_urls_from_env() -> None: del os.environ["CODEGATE_PROVIDER_VLLM_URL"] +def test_certificate_config_from_file(tmp_path: Path) -> None: + """Test loading certificate configuration from file.""" + config_file = tmp_path / "config.yaml" + cert_config = { + "certs_dir": "./custom-certs", + "ca_cert": "custom-ca.crt", + "ca_key": "custom-ca.key", + "server_cert": "custom-server.crt", + "server_key": "custom-server.key", + } + with open(config_file, "w") as f: + yaml.dump(cert_config, f) + + config = Config.from_file(config_file) + assert config.certs_dir == cert_config["certs_dir"] + assert config.ca_cert == cert_config["ca_cert"] + assert config.ca_key == cert_config["ca_key"] + assert config.server_cert == cert_config["server_cert"] + assert config.server_key == cert_config["server_key"] + + +def test_certificate_config_from_env() -> None: + """Test loading certificate configuration from environment variables.""" + os.environ.update({ + "CODEGATE_CERTS_DIR": "./env-certs", + "CODEGATE_CA_CERT": "env-ca.crt", + "CODEGATE_CA_KEY": "env-ca.key", + "CODEGATE_SERVER_CERT": "env-server.crt", + "CODEGATE_SERVER_KEY": "env-server.key", + }) + try: + config = Config.from_env() + assert config.certs_dir == "./env-certs" + assert config.ca_cert == "env-ca.crt" + assert config.ca_key == "env-ca.key" + assert config.server_cert == "env-server.crt" + assert config.server_key == "env-server.key" + finally: + for key in [ + "CODEGATE_CERTS_DIR", + "CODEGATE_CA_CERT", + "CODEGATE_CA_KEY", + "CODEGATE_SERVER_CERT", + "CODEGATE_SERVER_KEY", + ]: + os.environ.pop(key, None) + + def test_invalid_log_level() -> None: """Test handling of invalid log level.""" with pytest.raises(ConfigurationError): diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 00000000..0e4db9f8 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,150 @@ +"""Tests for the server module.""" + +from unittest.mock import MagicMock, patch + +import pytest +from fastapi.testclient import TestClient +from fastapi.middleware.cors import CORSMiddleware + +from codegate import __version__ +from codegate.server import init_app +from codegate.providers.registry import ProviderRegistry +from codegate.pipeline.secrets.manager import SecretsManager + + +@pytest.fixture +def mock_secrets_manager(): + """Create a mock secrets manager.""" + return MagicMock(spec=SecretsManager) + + +@pytest.fixture +def mock_provider_registry(): + """Create a mock provider registry.""" + return MagicMock(spec=ProviderRegistry) + + +@pytest.fixture +def test_client() -> TestClient: + """Create a test client for the FastAPI application.""" + app = init_app() + return TestClient(app) + + +def test_app_initialization() -> None: + """Test that the FastAPI application initializes correctly.""" + app = init_app() + assert app is not None + assert app.title == "CodeGate" + assert app.version == __version__ + + +def test_cors_middleware() -> None: + """Test that CORS middleware is properly configured.""" + app = init_app() + cors_middleware = next( + m for m in app.middleware if isinstance(m, CORSMiddleware) + ) + assert cors_middleware.options.allow_origins == ["*"] + assert cors_middleware.options.allow_credentials is True + assert cors_middleware.options.allow_methods == ["*"] + assert cors_middleware.options.allow_headers == ["*"] + + +def test_health_check(test_client: TestClient) -> None: + """Test the health check endpoint.""" + response = test_client.get("/health") + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +@patch("codegate.server.ProviderRegistry") +@patch("codegate.server.SecretsManager") +def test_provider_registration(mock_secrets_mgr, mock_registry) -> None: + """Test that all providers are registered correctly.""" + init_app() + + # Verify SecretsManager was initialized + mock_secrets_mgr.assert_called_once() + + # Verify ProviderRegistry was initialized with the app + mock_registry.assert_called_once() + + # Verify all providers were registered + registry_instance = mock_registry.return_value + assert registry_instance.add_provider.call_count == 5 # openai, anthropic, llamacpp, vllm, ollama + + # Verify specific providers were registered + provider_names = [ + call.args[0] for call in registry_instance.add_provider.call_args_list + ] + assert "openai" in provider_names + assert "anthropic" in provider_names + assert "llamacpp" in provider_names + assert "vllm" in provider_names + assert "ollama" in provider_names + + +@patch("codegate.server.CodegateSignatures") +def test_signatures_initialization(mock_signatures) -> None: + """Test that signatures are initialized correctly.""" + init_app() + mock_signatures.initialize.assert_called_once_with("signatures.yaml") + + +def test_pipeline_initialization() -> None: + """Test that pipelines are initialized correctly.""" + app = init_app() + + # Access the provider registry to check pipeline configuration + registry = next( + (route for route in app.routes if hasattr(route, "registry")), + None + ) + + if registry: + for provider in registry.registry.values(): + # Verify each provider has the required pipelines + assert hasattr(provider, "pipeline_processor") + assert hasattr(provider, "fim_pipeline_processor") + assert hasattr(provider, "output_pipeline_processor") + + +def test_dashboard_routes() -> None: + """Test that dashboard routes are included.""" + app = init_app() + routes = [route.path for route in app.routes] + + # Verify dashboard endpoints are included + dashboard_routes = [route for route in routes if route.startswith("/dashboard")] + assert len(dashboard_routes) > 0 + + +def test_system_routes() -> None: + """Test that system routes are included.""" + app = init_app() + routes = [route.path for route in app.routes] + + # Verify system endpoints are included + assert "/health" in routes + + +@pytest.mark.asyncio +async def test_async_health_check() -> None: + """Test the health check endpoint with async client.""" + app = init_app() + + async with TestClient(app) as ac: + response = await ac.get("/health") + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +def test_error_handling(test_client: TestClient) -> None: + """Test error handling for non-existent endpoints.""" + response = test_client.get("/nonexistent") + assert response.status_code == 404 + + # Test method not allowed + response = test_client.post("/health") # Health endpoint only allows GET + assert response.status_code == 405