ctf-ad-mirror-tool-rust/src/flags/mod.rs

//! Optional format-preserving tokenization with an authenticated encrypted vault.
//! This is not client authorization or encryption of the transport itself.
mod stream;
mod vault;
use anyhow::{Result, ensure};
use serde::{Deserialize, Serialize};
use std::{collections::HashSet, path::PathBuf};
pub use stream::Transformer;
pub use vault::{FlagStatus, Vault};
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct FlagConfig {
    pub database: PathBuf,
    pub key_env: String,
    pub prefix: String,
    pub alphabet: String,
    pub payload_length: usize,
    pub suffix: String,
    pub ttl_seconds: u64,
    pub max_entries: usize,
    pub partial_timeout_ms: u64,
}
impl Default for FlagConfig {
    fn default() -> Self {
        Self {
            database: "data/flags.sqlite3".into(),
            key_env: "MIRROR_FLAG_KEY".into(),
            prefix: String::new(),
            alphabet: "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789".into(),
            payload_length: 31,
            suffix: "=".into(),
            ttl_seconds: 86400,
            max_entries: 10000,
            partial_timeout_ms: 5000,
        }
    }
}
impl FlagConfig {
    pub fn validate(&self) -> Result<()> {
        ensure!(
            !self.database.as_os_str().is_empty(),
            "flags.database is required"
        );
        ensure!(
            !self.key_env.is_empty()
                && self
                    .key_env
                    .bytes()
                    .all(|c| c.is_ascii_alphanumeric() || c == b'_'),
            "invalid flags.key_env"
        );
        ensure!(
            (16..=128).contains(&self.payload_length),
            "flags.payload_length must be 16..128"
        );
        let alphabet: HashSet<_> = self.alphabet.bytes().collect();
        ensure!(
            (16..=94).contains(&alphabet.len())
                && alphabet.len() == self.alphabet.len()
                && self.alphabet.bytes().all(|c| c.is_ascii_graphic()),
            "flags.alphabet must contain 16..94 distinct printable ASCII bytes"
        );
        ensure!(
            self.prefix.len() <= 32
                && self.suffix.len() <= 32
                && self
                    .prefix
                    .bytes()
                    .chain(self.suffix.bytes())
                    .all(|c| c.is_ascii_graphic()),
            "flag prefix/suffix must be printable ASCII, at most 32 bytes each"
        );
        ensure!(
            (60..=604800).contains(&self.ttl_seconds),
            "flags.ttl_seconds must be 60..604800"
        );
        ensure!(
            (1..=100000).contains(&self.max_entries),
            "flags.max_entries must be 1..100000"
        );
        ensure!(
            (100..=60000).contains(&self.partial_timeout_ms),
            "flags.partial_timeout_ms must be 100..60000"
        );
        Ok(())
    }
    pub fn width(&self) -> usize {
        self.prefix.len() + self.payload_length + self.suffix.len()
    }
    pub fn matches_prefix(&self, bytes: &[u8]) -> bool {
        bytes.len() <= self.width()
            && bytes.iter().enumerate().all(|(i, &c)| {
                if i < self.prefix.len() {
                    c == self.prefix.as_bytes()[i]
                } else if i < self.prefix.len() + self.payload_length {
                    self.alphabet.as_bytes().contains(&c)
                } else {
                    c == self.suffix.as_bytes()[i - self.prefix.len() - self.payload_length]
                }
            })
    }
    pub(crate) fn format_id(&self) -> Vec<u8> {
        serde_json::to_vec(&(
            &self.prefix,
            &self.alphabet,
            self.payload_length,
            &self.suffix,
        ))
        .expect("format serialization")
    }
}