iocaine 2.2.0

The deadliest poison known to AI
Documentation
// SPDX-FileCopyrightText: 2025 Gergely Nagy
// SPDX-FileContributor: Gergely Nagy
//
// SPDX-License-Identifier: MIT

#[cfg(all(not(target_env = "musl"), feature = "jemalloc"))]
#[global_allocator]
static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

use anyhow::Result;
use axum::http::header::{self, HeaderMap, HeaderName};
use clap::Parser;
use std::error::Error;

use iocaine::means_of_production::MeansOfProduction;

#[derive(Debug, Parser)]
#[command(version, about)]
pub struct Args {
    #[arg(short = 's', long)]
    pub script_path: String,

    #[arg(short = 'P', long)]
    pub path: Option<String>,
    #[arg(short = 'A', long, default_value_t = Self::default_user_agent())]
    pub user_agent: String,
    #[arg(short = 'H', long = "header", value_parser = Self::parse_key_val::<String, String>)]
    pub headers: Vec<(String, String)>,
}

impl Args {
    fn default_user_agent() -> String {
        "Mozilla/5.0 (X11; Linux x86_64; rv:137.0) Gecko/20100101 Firefox/137.0".to_owned()
    }

    fn parse_key_val<T, U>(
        s: &str,
    ) -> std::result::Result<(T, U), Box<dyn Error + Send + Sync + 'static>>
    where
        T: std::str::FromStr,
        T::Err: Error + Send + Sync + 'static,
        U: std::str::FromStr,
        U::Err: Error + Send + Sync + 'static,
    {
        let pos = s
            .find('=')
            .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{s}`"))?;
        Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
    }
}

fn main() -> Result<()> {
    let args = Args::parse();

    let handler = MeansOfProduction::new(&args.script_path)?;

    let mut headers = HeaderMap::new();
    headers.insert(header::HOST, "example.com".parse()?);
    headers.insert(header::USER_AGENT, args.user_agent.parse()?);

    for (key, value) in args.headers {
        headers.insert(HeaderName::from_bytes(key.as_bytes())?, value.parse()?);
    }

    let verdict = handler.decide(headers, args.path, "GET");
    println!("{verdict:?}");

    Ok(())
}