server.py

import ssl
import sys
from xmlrpc.server import SimpleXMLRPCServer, SimpleXMLRPCRequestHandler
from socketserver import ThreadingMixIn
import threading
import time
import os
import typer
from loguru import logger

from utils import generate_self_signed_cert, Data, convert_value_from_xmlrpc

# Ensure logs directory exists
logs_dir = "logs"
os.makedirs(logs_dir, exist_ok=True)

# Configure loguru logger. Loguru installs a DEBUG-level stderr sink by
# default, which floods the console with internal tracing on every request;
# replace it so both sinks honour LOG_LEVEL (INFO unless overridden).
LOG_LEVEL = os.environ.get("LOG_LEVEL", "INFO").upper()
logger.remove()
logger.add(sys.stderr, level=LOG_LEVEL)
logger.add(os.path.join(logs_dir, "server.log"), rotation="1 MB", level=LOG_LEVEL)

app = typer.Typer(help="XML-RPC Server - Start HTTP and/or HTTPS servers")

# How long a connection may stay idle mid-request before it is dropped.
REQUEST_TIMEOUT = 30.0
# Generous: the HTTPS server may need to generate a certificate first.
STARTUP_TIMEOUT = 60.0


class BoundedXMLRPCRequestHandler(SimpleXMLRPCRequestHandler):
    """Request handler that refuses to wait forever on a silent client.

    Without a timeout a client that opens a connection and never sends a
    complete request pins a worker thread indefinitely.
    """
    timeout = REQUEST_TIMEOUT


class ThreadedXMLRPCServer(ThreadingMixIn, SimpleXMLRPCServer):
    """Threaded XML-RPC server with custom dispatch for kwargs support"""
    daemon_threads = True
    allow_reuse_address = True

    def __init__(self, addr, *args, ssl_context=None, **kwargs):
        self.ssl_context = ssl_context
        kwargs.setdefault("requestHandler", BoundedXMLRPCRequestHandler)
        kwargs.setdefault("allow_none", True)
        kwargs.setdefault("use_builtin_types", True)
        kwargs.setdefault("logRequests", False)
        super().__init__(addr, *args, **kwargs)
        logger.info("Initialized XML-RPC server on {}", addr)

    def get_request(self):
        """Accept a connection, deferring any TLS handshake to the worker thread.

        Wrapping the *listening* socket makes SSLSocket.accept() run the
        handshake inside the single accept loop, so one client that connects
        and stays silent blocks every other client. Wrapping the accepted
        connection with do_handshake_on_connect=False moves the handshake to
        the first read, which happens on the per-request thread.
        """
        sock, addr = super().get_request()
        if self.ssl_context is not None:
            sock = self.ssl_context.wrap_socket(
                sock, server_side=True, do_handshake_on_connect=False
            )
        return sock, addr

    def handle_error(self, request, client_address):
        """Log connection-level failures without dumping a traceback to stderr."""
        exc = sys.exc_info()[1]
        if isinstance(exc, (ssl.SSLError, ConnectionError, TimeoutError)):
            logger.warning("Dropped connection from {}: {}", client_address, exc)
            return
        logger.opt(exception=True).error("Error processing request from {}", client_address)

    def _dispatch(self, method, params):
        """Custom dispatch method to handle Data objects for args and kwargs support"""
        try:
            # Handle introspection methods specially - don't wrap their results
            if method.startswith('system.'):
                return super()._dispatch(method, params)

            # Never expose private/dunder attributes over RPC. SimpleXMLRPCServer
            # blocks these in resolve_dotted_attribute; overriding _dispatch
            # bypasses that guard, which would otherwise make __dir__, __repr__
            # and friends remotely callable.
            if method.startswith('_'):
                error_msg = f'Method "{method}" is not supported!'
                logger.warning("Rejected private method call: {}", method)
                return Data(response_code=404, error=error_msg).__dict__

            # Get the function
            func = self.funcs.get(method)
            if func is None:
                # Try to get from instance
                if self.instance is not None:
                    func = getattr(self.instance, method, None)

            if func is None or not callable(func):
                error_msg = f'Method "{method}" is not supported!'
                logger.error(error_msg)
                return Data(response_code=404, error=error_msg).__dict__

            # Process parameters to extract Data objects and convert types
            args = []
            kwargs = {}

            for param in params:
                # Check if this is a Data object (serialized as dict with _is_data_object marker)
                if isinstance(param, dict) and param.get('_is_data_object', False):
                    # This is a Data object - extract args and kwargs and convert types
                    data_args = param.get('args', ())
                    data_kwargs = param.get('kwargs', {})
                    # Convert string representations back to int/float
                    converted_args = [convert_value_from_xmlrpc(arg) for arg in data_args]
                    converted_kwargs = {k: convert_value_from_xmlrpc(v) for k, v in data_kwargs.items()}
                    args.extend(converted_args)
                    kwargs.update(converted_kwargs)
                    logger.debug("Unpacked Data object: args={}, kwargs={}", converted_args, converted_kwargs)
                else:
                    # Regular parameter - convert if needed
                    converted_param = convert_value_from_xmlrpc(param)
                    args.append(converted_param)

            # Call function with unpacked args and kwargs
            logger.debug("Calling {} with args={}, kwargs={}", method, args, kwargs)
            result = func(*args, **kwargs)

            # Always wrap the result in a successful Data object with result attribute
            # Convert int/float results to strings for XML-RPC transmission
            return Data(response_code=200, result=result).__dict__

        except Exception as e:
            error_msg = str(e)
            logger.error("Error in dispatch for method {}: {}", method, error_msg)
            return Data(response_code=400, error=error_msg).__dict__

class MathFunctions:
    """Simple math functions for XML-RPC testing"""

    def add(self, x, y):
        result = x + y
        logger.debug("add({}, {}) = {}", x, y, result)
        return result

    def subtract(self, x, y):
        result = x - y
        logger.debug("subtract({}, {}) = {}", x, y, result)
        return result

    def multiply(self, x, y):
        result = x * y
        logger.debug("multiply({}, {}) = {}", x, y, result)
        return result

    def divide(self, x, y):
        if y == 0:
            logger.warning("Division by zero attempted: {} / {}", x, y)
            raise ValueError("Cannot divide by zero")
        result = x / y
        logger.debug("divide({}, {}) = {}", x, y, result)
        return result

def register_functions(server):
    """Register functions with the XML-RPC server"""
    logger.info("Registering XML-RPC functions")

    # Register math functions
    math_functions = MathFunctions()
    server.register_instance(math_functions)

    # Enable introspection functions
    server.register_introspection_functions()

    logger.info("Successfully registered math functions and introspection")

def start_http_server(host="localhost", port=8000, ready=None):
    """Start HTTP XML-RPC server

    Args:
        ready: optional threading.Event set once the server is listening.
    """
    try:
        logger.info("Starting HTTP XML-RPC server on {}:{}", host, port)

        server = ThreadedXMLRPCServer((host, port))
        register_functions(server)

        logger.success("HTTP XML-RPC server started successfully on {}:{}", host, port)
        logger.info("Server is ready to accept connections")

        if ready is not None:
            ready.set()

        server.serve_forever()
    except OSError as e:
        logger.error("Failed to start HTTP server on {}:{}: {}", host, port, e)
        raise
    except KeyboardInterrupt:
        logger.info("HTTP server shutdown requested")
    except Exception as e:
        logger.error("Unexpected error in HTTP server: {}", e)
        raise
    finally:
        logger.info("HTTP server stopped")

def start_https_server(host="localhost", port=8443, ready=None):
    """Start HTTPS XML-RPC server

    Args:
        ready: optional threading.Event set once the server is listening.
    """
    try:
        logger.info("Starting HTTPS XML-RPC server on {}:{}", host, port)

        # Generate self-signed certificate if it doesn't exist
        cert_file = "server.crt"
        key_file = "server.key"

        if not os.path.exists(cert_file) or not os.path.exists(key_file):
            logger.info("SSL certificate not found, generating self-signed certificate")
            generate_self_signed_cert(cert_file, key_file)
            logger.success("Self-signed certificate generated successfully")
        else:
            logger.info("Using existing SSL certificate")

        # Create SSL context
        context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        context.load_cert_chain(cert_file, key_file)

        # The context is handed to the server so each accepted connection is
        # wrapped on its own worker thread (see ThreadedXMLRPCServer.get_request).
        server = ThreadedXMLRPCServer((host, port), ssl_context=context)
        register_functions(server)

        logger.success("HTTPS XML-RPC server started successfully on {}:{}", host, port)
        logger.info("Server is ready to accept secure connections")

        if ready is not None:
            ready.set()

        server.serve_forever()
    except ssl.SSLError as e:
        logger.error("SSL configuration error: {}", e)
        raise
    except OSError as e:
        logger.error("Failed to start HTTPS server on {}:{}: {}", host, port, e)
        raise
    except KeyboardInterrupt:
        logger.info("HTTPS server shutdown requested")
    except Exception as e:
        logger.error("Unexpected error in HTTPS server: {}", e)
        raise
    finally:
        logger.info("HTTPS server stopped")

def _wait_until_ready(name, thread, ready, timeout=STARTUP_TIMEOUT):
    """Wait for a server thread to signal readiness.

    Returns False if the thread died during startup (e.g. the port was already
    in use) or failed to come up in time, so the caller can report a real
    failure instead of claiming success.
    """
    deadline = time.monotonic() + timeout
    while True:
        if ready.wait(0.1):
            return True
        if not thread.is_alive():
            logger.error("{} server thread exited before becoming ready", name)
            return False
        if time.monotonic() > deadline:
            logger.error("{} server did not become ready within {}s", name, timeout)
            return False

@app.command()
def http(host: str = "localhost", port: int = 8000):
    """Start HTTP XML-RPC server"""
    logger.info("HTTP server command invoked with host={}, port={}", host, port)
    start_http_server(host, port)

@app.command()
def https(host: str = "localhost", port: int = 8443):
    """Start HTTPS XML-RPC server"""
    logger.info("HTTPS server command invoked with host={}, port={}", host, port)
    start_https_server(host, port)

@app.command()
def both(
    http_host: str = "localhost",
    http_port: int = 8000,
    https_host: str = "localhost",
    https_port: int = 8443
):
    """Start both HTTP and HTTPS servers"""
    logger.info("Starting both servers - HTTP on {}:{}, HTTPS on {}:{}", http_host, http_port, https_host, https_port)

    # Start both servers in separate threads
    http_ready = threading.Event()
    https_ready = threading.Event()
    http_thread = threading.Thread(
        target=start_http_server, args=(http_host, http_port, http_ready),
        daemon=True, name="http-server",
    )
    https_thread = threading.Thread(
        target=start_https_server, args=(https_host, https_port, https_ready),
        daemon=True, name="https-server",
    )

    http_thread.start()
    https_thread.start()

    # Confirm both are actually listening before reporting success. Wait on
    # both regardless so a single run reports every failure.
    http_ok = _wait_until_ready("HTTP", http_thread, http_ready)
    https_ok = _wait_until_ready("HTTPS", https_thread, https_ready)
    if not (http_ok and https_ok):
        logger.error("Not all servers started; exiting")
        print("Failed to start all servers - see logs for details")
        raise typer.Exit(1)

    logger.success("Both HTTP and HTTPS servers started successfully")
    logger.info("Press Ctrl+C to stop both servers")

    try:
        while True:
            # Daemon threads die silently; without this the process would keep
            # running while serving nothing.
            for name, thread in (("HTTP", http_thread), ("HTTPS", https_thread)):
                if not thread.is_alive():
                    logger.error("{} server thread stopped unexpectedly", name)
                    print(f"{name} server stopped unexpectedly - see logs for details")
                    raise typer.Exit(1)
            time.sleep(1)
    except KeyboardInterrupt:
        logger.info("Shutdown signal received for both servers")
        print("\nShutting down servers...")
    finally:
        logger.info("Both servers stopped")

if __name__ == "__main__":
    with logger.catch():
        typer.echo("Starting XML-RPC server application...")
        app()