continuum-proxy/src/printer/bambu.rs

use std::net::{SocketAddr, TcpStream, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::time::Duration;

use super::{GenericPrinter, PrinterBase, PrinterError};

// region: Certificates

/// The canonical shape for every Bambu MQTT message lives in
/// continuum-schemas, not here - see `send_gcode_line` below, which builds
/// an instance of it directly via `continuum-schema-instance` instead of a
/// hand-written serde struct. `v1`/`v2` underneath this directory are
/// BambuMQTT's two message schemas, not to be confused with the MQTT
/// *transport* version (v3.1.1 vs v5) below - they happen to move
/// together on real Bambu firmware, but they're two different things.
const BAMBU_SCHEMA_DIR: &str = concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../continuum-schemas/schemas/bambulab"
);

fn get_schema_for_command(version: BambuProtocolVersion, root: String, command: String) -> PathBuf {
    Path::new(BAMBU_SCHEMA_DIR)
        .join(version.schema_dir())
        .join(format!("request/{root}/{command}.schema.json"))
}

/// Bambu's shared LAN-mode root CA ("BBL CA2 RSA") - see certs/README.md
/// for how this was obtained and verified. `include_bytes!` embeds it into
/// the compiled binary at build time, so there's no file to ship alongside
/// the executable.
const BUNDLED_CA: &[u8] = include_bytes!("../../certs/bambu_ca2.pem");

/// Which certificate to trust when connecting to a Bambu printer's MQTTS
/// port (8883 - Bambu's LAN mode is TLS-only, there's no unencrypted
/// fallback).
pub enum BambuTls {
    /// Bambu's shared root CA. Correct for current-generation printers.
    BundledCa,
    /// A specific certificate instead - for a printer whose firmware
    /// doesn't chain to the shared CA (P1P, at least).
    Custom(Vec<u8>),
    /// Skip certificate verification entirely. The least safe option, only
    /// reasonable because LAN mode never leaves your local network.
    Insecure,
}

impl BambuTls {
    /// Resolves the trust mode for *one* printer. Call this per printer,
    /// passing `config.bambu_ca_cert_overrides.get(printer_id)` - there's
    /// no single override for the whole fleet, because not every Bambu
    /// model chains to the same CA. `require_valid_cert = false` always
    /// means `Insecure`, regardless of whether that printer also has an
    /// override entry.
    pub fn resolve(cert_path: Option<&Path>, require_valid_cert: bool) -> anyhow::Result<Self> {
        if !require_valid_cert {
            return Ok(BambuTls::Insecure);
        }
        match cert_path {
            Some(path) => Ok(BambuTls::Custom(std::fs::read(path)?)),
            None => Ok(BambuTls::BundledCa),
        }
    }

    /// The CA certificate bytes to hand to the TLS layer - empty when
    /// `Insecure`, since there's nothing to verify against.
    pub fn ca_bytes(&self) -> &[u8] {
        match self {
            BambuTls::BundledCa => BUNDLED_CA,
            BambuTls::Custom(bytes) => bytes,
            BambuTls::Insecure => &[],
        }
    }

    /// Builds a `native_tls` connector configured for this trust mode:
    /// trusting only `ca_bytes()`, or accepting anything for `Insecure`.
    /// `pub(crate)`, not private: `bambu_mqtt.rs` (a sibling module, not a
    /// descendant of this one) reuses this exact connector to configure
    /// rumqttc's TLS transport, instead of every consumer's TLS trust
    /// decisions living in two separate places.
    pub(crate) fn build_connector(&self) -> anyhow::Result<native_tls::TlsConnector> {
        let mut builder = native_tls::TlsConnector::builder();

        // Bambu printer certificates don't carry their LAN IP as a Subject
        // Alternative Name - it's DHCP-assigned, so baking it in at
        // manufacture time wouldn't make sense. Hostname verification has
        // to be off regardless of trust mode: what matters is whether the
        // certificate chains to a CA we trust, not whether its SAN matches
        // the specific IP we happened to dial today. Chain verification
        // itself stays fully enforced below for BundledCa/Custom - this
        // narrows *what* gets checked, it doesn't turn checking off.
        builder.danger_accept_invalid_hostnames(true);

        if matches!(self, BambuTls::Insecure) {
            builder.danger_accept_invalid_certs(true);
        } else {
            builder.add_root_certificate(native_tls::Certificate::from_pem(self.ca_bytes())?);
        }

        Ok(builder.build()?)
    }
}

fn resolve_addr(host: &str) -> anyhow::Result<SocketAddr> {
    (host, 8883)
        .to_socket_addrs()?
        .next()
        .ok_or_else(|| anyhow::anyhow!("could not resolve {host}"))
}

/// Minimal DER -> PEM encoder (base64, 64-char lines, standard
/// header/footer) - just enough to save a fetched certificate in the same
/// format as the bundled one.
fn der_to_pem(der: &[u8]) -> Vec<u8> {
    use base64::Engine;
    let encoded = base64::engine::general_purpose::STANDARD.encode(der);

    let mut pem = String::from("-----BEGIN CERTIFICATE-----\n");
    for line in encoded.as_bytes().chunks(64) {
        pem.push_str(std::str::from_utf8(line).expect("base64 output is always ASCII"));
        pem.push('\n');
    }
    pem.push_str("-----END CERTIFICATE-----\n");
    pem.into_bytes()
}

/// Re-encodes a PEM certificate to DER, so it can be byte-compared against
/// what a live handshake presents (`Certificate::to_der()`).
fn pem_to_der(pem: &[u8]) -> anyhow::Result<Vec<u8>> {
    Ok(native_tls::Certificate::from_pem(pem)?.to_der()?)
}

// endregion: Certificates

// region: BambuMQTTProtocol
/// Which BambuMQTT message schema (and, on real firmware, which MQTT
/// transport version) a connection ended up using. This is an *outcome*
/// `connect()` settles on, not something you can assume from the printer
/// model - a firmware update can move the same physical X1C from v1 to
/// v2. Every later message this printer sends is built against
/// `schemas/bambulab/{v1,v2}/` depending on which one this is.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BambuProtocolVersion {
    /// MQTT v3.1.1, BambuMQTT v1 - older firmware (P1P/P1S, original X1 launch firmware).
    V1,
    /// MQTT v5, BambuMQTT v2 - current firmware.
    V2,
}

impl BambuProtocolVersion {
    fn schema_dir(self) -> &'static str {
        match self {
            BambuProtocolVersion::V1 => "v1",
            BambuProtocolVersion::V2 => "v2",
        }
    }
}

/// How to pick the protocol generation for a given printer - config
/// *input*, as opposed to `BambuProtocolVersion` (the outcome `connect()`
/// settles on). `Auto` is the default for exactly the reason
/// `BambuProtocolVersion` isn't derivable from config alone: you can't
/// reliably know which one a specific printer speaks today without
/// asking it. `ForceV1`/`ForceV2` stay available as an override - useful
/// if auto-detection ever misbehaves, or to test one path deliberately.
#[derive(Debug, Clone, Copy)]
pub enum BambuProtocolPreference {
    /// Try MQTT v5 first; if the broker's CONNACK rejects that protocol
    /// version (older firmware only understands v3.1.1), fall back and
    /// reconnect with v3.1.1 instead.
    Auto,
    ForceV1,
    ForceV2,
}
// endregion: BambuMQTTProtocol

/// A Bambu printer. Only one struct now, not two: the old
/// `BambuGenericPrinter` (shared fields/behavior) + `BambuV1Printer`/
/// `BambuV2Printer` (thin per-version wrappers) split existed to let two
/// *separate types* share a base - but since which protocol version a
/// printer uses is now something `connect()` discovers at runtime rather
/// than something config picks between two types for, there's only ever
/// one type to begin with, so the extra layer stopped earning its keep.
pub struct BambuPrinter {
    base: PrinterBase,
    pub access_code: String,
    pub sn: String,
    tls: BambuTls,
    preference: BambuProtocolPreference,
    /// `None` until `connect()` succeeds. Carries the live MQTT session
    /// alongside the protocol version it settled on, since the two are
    /// decided together and always used together - see `bambu_mqtt.rs`.
    connection: Option<(super::bambu_mqtt::BambuMqttSession, BambuProtocolVersion)>,
}

impl BambuPrinter {
    pub fn new(
        id: &str,
        name: &str,
        host: &str,
        access_code: &str,
        sn: &str,
        tls: BambuTls,
        preference: BambuProtocolPreference,
    ) -> Self {
        Self {
            base: PrinterBase {
                id: id.into(),
                name: name.into(),
                host: host.into(),
            },
            access_code: access_code.into(),
            sn: sn.into(),
            tls,
            preference,
            connection: None,
        }
    }

    pub fn base(&self) -> &PrinterBase {
        &self.base
    }

    /* A real (blocking, one-shot) TLS handshake to port 8883 - no MQTT
    /// protocol, just "does the certificate verify" against this printer's
    /// configured trust mode.
    ///
    /// `BundledCa` verifies normally (Bambu's shared CA is a proper,
    /// well-formed root - OpenSSL's usual chain validation handles it
    /// fine). `Custom` does something different: *fingerprint pinning*,
    /// not chain validation. A device's own leaf certificate often isn't a
    /// well-formed CA (no `CA:TRUE`), and the device may present it
    /// alongside a separate self-signed root you never captured - asking
    /// OpenSSL's normal path-building to accept an arbitrary leaf as a
    /// trust anchor is unreliable for exactly that reason (this is a real
    /// bug that shipped and failed against a real P1P: `add_root_certificate`
    /// with a captured leaf produced `self-signed certificate in
    /// certificate chain` / `unable to get local issuer certificate` on
    /// the very next connection). So instead: connect with verification
    /// off, then compare the certificate actually presented against the
    /// exact bytes pinned, byte for byte. A mismatch is a hard failure -
    /// no PKI judgment call, just "is this the same certificate as before." */
    pub fn test_tls_handshake(&self) -> anyhow::Result<()> {
        let addr = resolve_addr(&self.base.host)?;
        let stream = TcpStream::connect_timeout(&addr, Duration::from_secs(5))?;

        if let BambuTls::Custom(pinned_pem) = &self.tls {
            let tls_stream = BambuTls::Insecure
                .build_connector()?
                .connect(&self.base.host, stream)?;
            let presented = tls_stream
                .peer_certificate()?
                .ok_or_else(|| anyhow::anyhow!("{} presented no certificate", self.base.name))?
                .to_der()?;

            if presented != pem_to_der(pinned_pem)? {
                anyhow::bail!(
                    "{} presented a certificate different from the one pinned - could be a legitimate \
                     certificate rotation, could be something worse; re-run fetch_bambu_cert deliberately \
                     if you're sure it's the former",
                    self.base.name
                );
            }
            return Ok(());
        }

        self.tls
            .build_connector()?
            .connect(&self.base.host, stream)?;
        Ok(())
    }

    /* "Trust on first connect": connects once with certificate
    /// verification disabled, captures whatever certificate the printer
    /// presents, and returns it as PEM bytes.
    ///
    /// This is meaningfully weaker than verifying against a CA you already
    /// trust - anyone on the network at the exact moment you run this gets
    /// trusted forever after. That's why it's a method you call once, by
    /// hand (see `examples/fetch_bambu_cert.rs`), rather than something
    /// `connect()` falls back to silently. Save the result to a file and
    /// point that printer's `ca_cert_path` in `printers.toml` at it. */
    pub fn fetch_certificate(&self) -> anyhow::Result<Vec<u8>> {
        let addr = resolve_addr(&self.base.host)?;
        let stream = TcpStream::connect_timeout(&addr, Duration::from_secs(5))?;
        let tls_stream = BambuTls::Insecure
            .build_connector()?
            .connect(&self.base.host, stream)?;

        let cert = tls_stream
            .peer_certificate()?
            .ok_or_else(|| anyhow::anyhow!("{} presented no certificate", self.base.name))?;

        Ok(der_to_pem(&cert.to_der()?))
    }

    /* Makes sure this printer has a *working* trust configuration,
    /// auto-pinning a certificate via trust-on-first-connect if needed -
    /// the "smooth, no manual steps" version of the fetch_bambu_cert
    /// workflow, built the way SSH handles host keys:
    ///
    /// - A pinned certificate already on disk (`cert_dir/<sn>.pem`) is used
    ///   directly. No new trust decision gets made on every run - that
    ///   already happened once, this just re-verifies against it. Keyed by
    ///   *serial number*, not `id`: `id` is just an arbitrary label you
    ///   picked in printers.toml - rename a printer's `id` and the pin
    ///   should still find it, because a pin certifies "this physical
    ///   printer", and the serial number is the one thing about it that
    ///   can't change.
    /// - No pin yet, and the bundled CA verifies fine (current-generation
    ///   printers): nothing to do.
    /// - No pin yet, and the bundled CA does *not* verify (P1P, etc.):
    ///   fetches the certificate the printer presents, confirms it
    ///   produces a working connection, and saves it to `cert_dir/<sn>.pem`
    ///   for every run after this one. This is the one moment a MITM
    ///   active on your network *right now* could plant a certificate that
    ///   gets trusted from then on - the same tradeoff SSH accepts on a
    ///   first connection.
    /// - A printer with an *explicit* pin (you set `ca_cert_path` yourself
    ///   in printers.toml) never auto-pins over it - a verification
    ///   failure there is a real error, not something to paper over,
    ///   since that's exactly the signal pinning exists to give you (the
    ///   cert rotated, or something worse). */
    pub fn ensure_trusted(&mut self, cert_dir: &Path) -> anyhow::Result<()> {
        let pinned_path = cert_dir.join(format!("{}.pem", self.sn));

        if matches!(self.tls, BambuTls::BundledCa) && pinned_path.exists() {
            self.tls = BambuTls::Custom(std::fs::read(&pinned_path)?);
        }

        match self.test_tls_handshake() {
            Ok(()) => Ok(()),
            Err(err) if matches!(self.tls, BambuTls::BundledCa) => {
                eprintln!(
                    "{}: bundled CA didn't verify ({err}); auto-pinning on first connect",
                    self.base.name
                );
                let pem = self.fetch_certificate()?;
                std::fs::create_dir_all(cert_dir)?;
                std::fs::write(&pinned_path, &pem)?;
                self.tls = BambuTls::Custom(pem);
                self.test_tls_handshake()
            }
            Err(err) => Err(err),
        }
    }
}

impl GenericPrinter for BambuPrinter {
    fn base(&self) -> &PrinterBase {
        &self.base
    }

    /// Real, not a stub: opens the MQTT session (see `bambu_mqtt.rs` for
    /// `preference`'s `Auto`-tries-v5-then-v3.1.1 fallback) and keeps it -
    /// same reasoning as `PrusaSerialPrinter::connect` holding its port
    /// open, just an MQTT session instead of a serial port.
    async fn connect(&mut self) -> Result<(), PrinterError> {
        let (session, version) = super::bambu_mqtt::BambuMqttSession::open(
            &self.base.host,
            &self.access_code,
            &self.sn,
            &self.tls,
            self.preference,
        )
        .await?;

        println!("[{}] connected over MQTT ({version:?})", self.base.name);
        self.connection = Some((session, version));
        Ok(())
    }

    /// Real: a raw pass-through publish to `device/{sn}/request`, no
    /// category envelope opinion - `send_gcode_line` below builds and
    /// sends its own properly-wrapped request instead of going through
    /// this.
    async fn publish(&mut self, payload: String) -> Result<(), PrinterError> {
        let (session, _) =
            self.connection.as_ref().ok_or_else(|| PrinterError(format!("{} is not connected", self.base.name)))?;
        session.publish_raw(payload).await
    }

    /// Same MQTT gcode-injection mechanism regardless of protocol
    /// generation - just routed through `send_gcode_line` below instead of
    /// its own separate stub, since that's genuinely the same operation.
    async fn set_fan_speed(&mut self, percent: u8) -> Result<(), PrinterError> {
        let pwm = super::percent_to_pwm(percent);
        self.send_gcode_line(format!("M106 P1 S{pwm}")).await
    }

    /// Real version FTPS-`PUT`s `data` to the printer's own `/cache/`
    /// directory (that's storage the printer's firmware owns, same idea
    /// as PrusaLink's SD card), then publishes a `project_file` MQTT
    /// command referencing it by name so the printer actually starts
    /// printing it - two separate steps, both stubbed here as one line
    /// for now.
    async fn upload_project_file(
        &mut self,
        file_name: String,
        data: Vec<u8>,
    ) -> Result<(), PrinterError> {
        println!(
            "[{}] FTPS upload {file_name} ({} bytes) to /cache/",
            self.base.name,
            data.len()
        );
        Ok(())
    }

    /// Builds the real JSON payload straight from continuum-schemas'
    /// `schemas/bambulab/{v1,v2}/request/print/gcode_line.schema.json` -
    /// which one depends on `self.version`, settled by `connect()` above -
    /// instead of a hand-written serde struct: `sequence_id` and the
    /// `command` const get filled in by the schema itself, `param` is the
    /// one field we actually set. Real, end to end: publishes over the
    /// live MQTT session `connect()` opened, no waiting for a reply -
    /// `gcode_line` is fire-and-forget, the printer just executes it.
    async fn send_gcode_line(&mut self, gcode_line: String) -> Result<(), PrinterError> {
        let (session, version) =
            self.connection.as_ref().ok_or_else(|| PrinterError(format!("{} is not connected", self.base.name)))?;

        let mut body = continuum_schema_instance::build_instance_from_file(&get_schema_for_command(
            *version,
            "print".into(),
            "gcode_line".into(),
        ))
        .map_err(|err| PrinterError(format!("failed to build gcode_line request: {err}")))?;
        body["param"] = serde_json::Value::String(gcode_line);

        session.publish("print", &body).await
    }
}