qssh 0.4.4

//! QSSH protocol messages

use serde::{Serialize, Deserialize};
use crate::{PqAlgorithm, KexAlgorithm};

/// The current supported version of the QSSH protocol
/// Represented as a (Major, Minor) tuple (e.g., 0.1)
pub const PROTOCOL_VERSION: (u8, u8) = (0, 1);

/// Protocol messages
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Message {
    /// Version exchange
    Version(VersionMessage),
    
    /// Client hello
    ClientHello(ClientHelloMessage),
    
    /// Server hello
    ServerHello(ServerHelloMessage),
    
    /// Key exchange
    KeyExchange(KeyExchangeMessage),
    
    /// Authentication
    Auth(AuthMessage),
    
    /// Channel operations
    Channel(ChannelMessage),
    
    /// Disconnect
    Disconnect(DisconnectMessage),
    
    /// Keepalive
    Ping(u64),
    Pong(u64),
    
    /// Key rotation
    Rekey(RekeyMessage),

    /// Global request (not channel-specific, e.g., remote port forwarding)
    GlobalRequest(GlobalRequestMessage),
    /// Global request success response
    GlobalRequestSuccess(GlobalRequestSuccessMessage),
    /// Global request failure
    GlobalRequestFailure,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VersionMessage {
    pub version: (u8, u8),
    pub software: String,
    pub comments: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClientHelloMessage {
    pub version: (u8, u8),
    pub random: [u8; 32],
    /// Supported key exchange algorithms (ordered by preference)
    pub kex_algorithms: Vec<KexAlgorithm>,
    pub sig_algorithms: Vec<PqAlgorithm>,
    pub ciphers: Vec<String>,
    pub qkd_capable: bool,
    pub extensions: Vec<Extension>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServerHelloMessage {
    pub version: (u8, u8),
    pub random: [u8; 32],
    /// Selected key exchange algorithm
    pub selected_kex: KexAlgorithm,
    pub selected_sig: PqAlgorithm,
    pub selected_cipher: String,
    /// Falcon public key (used for FalconSignedShares KEX)
    pub falcon_public_key: Vec<u8>,
    /// Key share for FalconSignedShares KEX
    pub key_share: Vec<u8>,
    /// Signature over key_share (for FalconSignedShares)
    pub key_share_signature: Vec<u8>,
    /// ML-KEM encapsulation key (used for MlKem768/MlKem1024 KEX)
    #[serde(default)]
    pub mlkem_encapsulation_key: Option<Vec<u8>>,
    /// X25519 public key (used for Hybrid KEX)
    #[serde(default)]
    pub x25519_public_key: Option<Vec<u8>>,
    pub qkd_endpoint: Option<String>,
    pub extensions: Vec<Extension>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyExchangeMessage {
    /// Falcon public key (for FalconSignedShares KEX)
    pub falcon_public_key: Vec<u8>,
    /// Key share for FalconSignedShares KEX
    pub key_share: Vec<u8>,
    /// Signature over key_share (for FalconSignedShares)
    pub key_share_signature: Vec<u8>,
    pub sphincs_public_key: Vec<u8>,
    /// ML-KEM ciphertext (used for MlKem768/MlKem1024 KEX)
    #[serde(default)]
    pub mlkem_ciphertext: Option<Vec<u8>>,
    /// X25519 public key from client (used for Hybrid KEX)
    #[serde(default)]
    pub x25519_public_key: Option<Vec<u8>>,
    pub qkd_proof: Option<Vec<u8>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AuthMessage {
    pub username: String,
    pub auth_method: AuthMethod,
    pub signature: Vec<u8>,
    pub session_id: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AuthMethod {
    PublicKey {
        algorithm: PqAlgorithm,
        public_key: Vec<u8>,
    },
    Password {
        password_hash: Vec<u8>,
    },
    Certificate {
        /// Serialized SshCertificate (bincode)
        certificate_data: Vec<u8>,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ChannelMessage {
    Open {
        channel_id: u32,
        channel_type: ChannelType,
        window_size: u32,
        max_packet_size: u32,
    },
    Accept {
        channel_id: u32,
        sender_channel: u32,
        window_size: u32,
        max_packet_size: u32,
    },
    Data {
        channel_id: u32,
        data: Vec<u8>,
    },
    WindowAdjust {
        channel_id: u32,
        bytes_to_add: u32,
    },
    Close {
        channel_id: u32,
    },
    Eof {
        channel_id: u32,
    },
    PtyRequest {
        channel_id: u32,
        term: String,
        width_chars: u32,
        height_chars: u32,
        width_pixels: u32,
        height_pixels: u32,
        modes: Vec<u8>,
    },
    ShellRequest {
        channel_id: u32,
    },
    ForwardRequest {
        channel_id: u32,
        remote_host: String,
        remote_port: u16,
    },
    ExecRequest {
        channel_id: u32,
        command: String,
    },
    X11Request {
        channel_id: u32,
        single_connection: bool,
        auth_protocol: String,
        auth_cookie: String,
        screen_number: u32,
    },
    SubsystemRequest {
        channel_id: u32,
        subsystem: String,
    },
    WindowChange {
        channel_id: u32,
        width_chars: u32,
        height_chars: u32,
        width_pixels: u32,
        height_pixels: u32,
    },
    /// Exit status from remote command (like SSH's exit-status request)
    ExitStatus {
        channel_id: u32,
        exit_code: u32,
    },
    /// Client requests agent forwarding on a session channel
    AgentForwardRequest {
        channel_id: u32,
    },
    /// Server acknowledges agent forwarding and provides the socket path
    AgentForwardSuccess {
        channel_id: u32,
        socket_path: String,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ChannelType {
    Session,
    DirectTcpIp,
    ForwardedTcpIp,
    X11,
    /// Agent forwarding channel (carries agent protocol messages)
    AgentForward,
    DirectTcpip {
        host: String,
        port: u16,
        originator_host: String,
        originator_port: u16,
    },
    ForwardedTcpip {
        connected_host: String,
        connected_port: u16,
        originator_host: String,
        originator_port: u16,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DisconnectMessage {
    pub reason_code: u32,
    pub description: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RekeyMessage {
    pub new_falcon_public_key: Vec<u8>,
    pub new_key_share: Vec<u8>,
    pub new_key_share_signature: Vec<u8>,
    pub request_qkd: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Extension {
    MaxPacketSize(u32),
    NoDelay(bool),
    QkdRequired(bool),
    Custom(String, Vec<u8>),
}

/// Global request message (not tied to a channel)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GlobalRequestMessage {
    pub request_type: GlobalRequestType,
    pub want_reply: bool,
}

/// Types of global requests
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GlobalRequestType {
    /// Request the server to listen on a port and forward connections back
    TcpipForward { bind_host: String, bind_port: u16 },
    /// Cancel a previously requested remote forward
    CancelTcpipForward { bind_host: String, bind_port: u16 },
}

/// Global request success response
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GlobalRequestSuccessMessage {
    /// Actual bound port (may differ from requested if client requested 0)
    pub bound_port: u16,
}

/// Disconnect reason codes
pub mod disconnect_reasons {
    /// The peer sent a message that violates the protocol state machine
    pub const PROTOCOL_ERROR: u32 = 1;
    /// The cryptographic handshake failed to verify.
    pub const KEY_EXCHANGE_FAILED: u32 = 2;
    /// The user provided invalid credentials or an untrusted certificate
    pub const AUTHENTICATION_FAILED: u32 = 3;
    pub const CONNECTION_LOST: u32 = 4;
    pub const BY_APPLICATION: u32 = 5;
    pub const TOO_MANY_CONNECTIONS: u32 = 6;
    pub const REKEY_REQUIRED: u32 = 7;
}