bambuddy/backend/tests/unit/services/mock_ftp_server.py

"""Mock implicit FTPS server for testing BambuFTPClient.

Built on pyftpdlib with implicit TLS support to match Bambu printer behavior.
Supports failure injection, custom AVBL command, and filesystem inspection.
"""

import logging
import os
import threading
import time

from pyftpdlib.authorizers import DummyAuthorizer
from pyftpdlib.handlers import TLS_FTPHandler
from pyftpdlib.servers import FTPServer


class ImplicitTLS_FTPHandler(TLS_FTPHandler):
    """FTP handler that wraps the socket in TLS before sending the 220 banner.

    This implements implicit FTPS (port 990 style) where the TLS handshake
    happens immediately on connect, before any FTP protocol exchange.
    pyftpdlib only natively supports explicit FTPS (AUTH TLS after connect).
    """

    # Per-class failure injection map: command -> (code, message, remaining_count)
    # -1 remaining_count = permanent failure
    _failure_map: dict = {}

    # AVBL command response (bytes available)
    _avbl_bytes: int = 1073741824  # 1 GB default

    # Register AVBL as a recognized FTP command (pyftpdlib requires this)
    proto_cmds = {
        **TLS_FTPHandler.proto_cmds,
        "AVBL": {
            "perm": None,
            "auth": True,
            "arg": None,
            "help": "Syntax: AVBL (get available bytes).",
        },
    }

    def handle(self):
        """Wrap socket in TLS immediately, then send 220 banner."""
        self.secure_connection(self.get_ssl_context())
        super().handle()

    def ftp_PROT(self, line):
        """Override PROT to auto-set _pbsz for implicit FTPS.

        In implicit FTPS the connection is already TLS-secured, so requiring
        a separate PBSZ command is unnecessary. Python's ftplib prot_c()
        doesn't send PBSZ first (unlike prot_p()), causing 503 errors.
        Real Bambu printers don't enforce this for implicit FTPS either.
        """
        self._pbsz = True
        return super().ftp_PROT(line)

    def _check_failure(self, command: str, line: str):
        """Check if a failure is injected for this command.

        Returns True if a failure response was sent, False otherwise.
        """
        if command in self._failure_map:
            code, message, remaining = self._failure_map[command]
            if remaining != 0:
                if remaining > 0:
                    self._failure_map[command] = (code, message, remaining - 1)
                    if remaining - 1 == 0:
                        del self._failure_map[command]
                self.respond(f"{code} {message}")
                return True
        return False

    def ftp_AVBL(self, line):
        """Handle custom AVBL command (available bytes on storage)."""
        self.respond(f"213 {self._avbl_bytes}")

    def ftp_RETR(self, file):
        if self._check_failure("RETR", file):
            return
        return super().ftp_RETR(file)

    def ftp_STOR(self, file):
        if self._check_failure("STOR", file):
            return
        return super().ftp_STOR(file)

    def ftp_DELE(self, line):
        if self._check_failure("DELE", line):
            return
        return super().ftp_DELE(line)

    def ftp_CWD(self, path):
        if self._check_failure("CWD", path):
            return
        return super().ftp_CWD(path)

    def ftp_LIST(self, path=""):
        if self._check_failure("LIST", path):
            return
        return super().ftp_LIST(path)

    def ftp_SIZE(self, path):
        if self._check_failure("SIZE", path):
            return
        # Override to allow SIZE in ASCII mode (real Bambu printers allow it,
        # and BambuFTPClient.get_file_size() doesn't set TYPE I first)
        if not self.fs.isfile(self.fs.realpath(path)):
            self.respond(f"550 {self.fs.fs2ftp(path)} is not retrievable.")
            return
        try:
            size = self.run_as_current_user(self.fs.getsize, path)
        except OSError as err:
            self.respond(f"550 {err}.")
        else:
            self.respond(f"213 {size}")

    def ftp_PASS(self, line):
        if self._check_failure("PASS", line):
            return
        return super().ftp_PASS(line)


class MockBambuFTPServer:
    """Manages a mock implicit FTPS server in a background thread.

    Simulates a Bambu printer FTP server with:
    - Implicit TLS (like real printers on port 990)
    - Standard Bambu directory structure
    - AVBL command support
    - Per-command failure injection for testing error paths
    """

    def __init__(
        self,
        host: str,
        port: int,
        root_dir: str,
        cert_path: str,
        key_path: str,
        access_code: str = "12345678",
    ):
        self.host = host
        self.port = port
        self.root_dir = root_dir
        self.cert_path = cert_path
        self.key_path = key_path
        self.access_code = access_code
        self._server: FTPServer | None = None
        self._thread: threading.Thread | None = None
        # Create a unique handler class per instance so _failure_map is isolated
        self._handler_class = type(
            "TestFTPHandler",
            (ImplicitTLS_FTPHandler,),
            {
                "_failure_map": {},
                "_avbl_bytes": 1073741824,
            },
        )

    def start(self):
        """Start the FTP server in a background daemon thread."""
        authorizer = DummyAuthorizer()
        authorizer.add_user("bblp", self.access_code, self.root_dir, perm="elradfmwMT")

        handler = self._handler_class
        handler.authorizer = authorizer
        handler.certfile = self.cert_path
        handler.keyfile = self.key_path
        handler.passive_ports = range(60000, 60101)
        handler.tls_control_required = False
        handler.tls_data_required = False
        # Reset ssl_context so it picks up our cert/key
        handler.ssl_context = None

        # Suppress pyftpdlib's noisy logging (startup/shutdown banners)
        # to avoid "I/O operation on closed file" errors when xdist
        # workers tear down while the daemon thread is still logging.
        logging.getLogger("pyftpdlib").setLevel(logging.CRITICAL)

        self._server = FTPServer((self.host, self.port), handler)
        self._server.max_cons = 10
        self._server.max_cons_per_ip = 5

        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        # Brief wait for server to be ready
        time.sleep(0.1)

    def stop(self):
        """Stop the FTP server and wait for thread to exit."""
        if self._server:
            self._server.close_all()
        if self._thread:
            self._thread.join(timeout=5)
        self._server = None
        self._thread = None

    def inject_failure(self, command: str, code: int, message: str, count: int = -1):
        """Inject a failure response for a specific FTP command.

        Args:
            command: FTP command name (RETR, STOR, DELE, CWD, LIST, SIZE, PASS)
            code: FTP response code (e.g. 550, 553)
            message: Response message
            count: Number of times to fail (-1 = permanent)
        """
        self._handler_class._failure_map[command] = (code, message, count)

    def clear_failures(self):
        """Remove all injected failures."""
        self._handler_class._failure_map.clear()

    def set_avbl_bytes(self, n: int):
        """Set the response value for the AVBL command."""
        self._handler_class._avbl_bytes = n

    def add_file(self, relative_path: str, content: bytes = b""):
        """Add a file to the server's filesystem."""
        full_path = os.path.join(self.root_dir, relative_path.lstrip("/"))
        os.makedirs(os.path.dirname(full_path), exist_ok=True)
        with open(full_path, "wb") as f:
            f.write(content)

    def add_directory(self, relative_path: str):
        """Create a directory in the server's filesystem."""
        full_path = os.path.join(self.root_dir, relative_path.lstrip("/"))
        os.makedirs(full_path, exist_ok=True)

    def file_exists(self, relative_path: str) -> bool:
        """Check if a file exists on the server."""
        full_path = os.path.join(self.root_dir, relative_path.lstrip("/"))
        return os.path.isfile(full_path)

    def read_file(self, relative_path: str) -> bytes:
        """Read file content from the server's filesystem."""
        full_path = os.path.join(self.root_dir, relative_path.lstrip("/"))
        with open(full_path, "rb") as f:
            return f.read()