tirith-core 0.3.0

use std::net::Ipv4Addr;

use crate::extract::ExtractedUrl;
use crate::threatdb::{self, Ecosystem, ThreatDb};
use crate::tokenize::{Segment, ShellType};
use crate::verdict::{Evidence, Finding, RuleId, Severity};

/// A reference to a package extracted from a shell command.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PackageRef {
    pub ecosystem: Ecosystem,
    pub name: String,
    pub version: Option<String>,
}

/// Split a `name<sep>version` string (e.g. `serde@1.0` or `rails:7.0`).
///
/// Returns `(name, Some(version))` when `sep` is found and the version part
/// is non-empty, otherwise `(input, None)`.
fn split_name_version(s: &str, sep: char) -> (&str, Option<String>) {
    if let Some(pos) = s.find(sep) {
        let name = &s[..pos];
        let ver = &s[pos + 1..];
        (
            name,
            if ver.is_empty() {
                None
            } else {
                Some(ver.to_string())
            },
        )
    } else {
        (s, None)
    }
}

/// Extract package references from tokenized shell segments.
///
/// Recognizes install/add commands for: pip, npm, yarn, pnpm, bun, npx,
/// cargo, gem, go, composer, dotnet.
///
/// Skips flags (tokens starting with `-`) and known non-package arguments
/// like `--index-url <url>`, `--save-dev`, etc.
pub fn extract_packages(segments: &[Segment]) -> Vec<PackageRef> {
    let mut packages = Vec::new();

    for seg in segments {
        let cmd = match &seg.command {
            Some(c) => c.to_lowercase(),
            None => continue,
        };

        // Strip path prefix: `/usr/bin/pip3` -> `pip3`.
        let cmd_name = cmd.rsplit('/').next().unwrap_or(&cmd);

        match cmd_name {
            "pip" | "pip3" | "uv" => {
                extract_pip_packages(&seg.args, &mut packages);
            }
            "npm" | "npx" | "yarn" | "pnpm" | "bun" => {
                extract_npm_packages(cmd_name, &seg.args, &mut packages);
            }
            "cargo" => {
                extract_cargo_packages(&seg.args, &mut packages);
            }
            "gem" => {
                extract_gem_packages(&seg.args, &mut packages);
            }
            "go" => {
                extract_go_packages(&seg.args, &mut packages);
            }
            "composer" => {
                extract_composer_packages(&seg.args, &mut packages);
            }
            "dotnet" => {
                extract_dotnet_packages(&seg.args, &mut packages);
            }
            "mvn" | "gradle" | "gradlew" => {
                extract_maven_packages(&seg.args, &mut packages);
            }
            _ => {}
        }
    }

    packages
}

/// Flags for pip that consume the next argument (so it should be skipped).
const PIP_ARG_FLAGS: &[&str] = &[
    "--index-url",
    "-i",
    "--extra-index-url",
    "--find-links",
    "-f",
    "--constraint",
    "-c",
    "--requirement",
    "-r",
    "--target",
    "-t",
    "--root",
    "--prefix",
    "--src",
    "--build",
    "-b",
    "--config-settings",
    "--global-option",
    "--install-option",
    "--proxy",
    "--retries",
    "--timeout",
    "--exists-action",
    "--trusted-host",
    "--cert",
    "--client-cert",
    "--cache-dir",
];

fn extract_pip_packages(args: &[String], packages: &mut Vec<PackageRef>) {
    let mut iter = args.iter();
    let mut found_install = false;
    while let Some(arg) = iter.next() {
        let lower = arg.to_lowercase();
        if !found_install {
            if lower == "install" {
                found_install = true;
            }
            continue;
        }

        if arg.starts_with('-') {
            if PIP_ARG_FLAGS.contains(&lower.as_str()) {
                let _ = iter.next();
            }
            continue;
        }

        // VCS / URL / local path installs aren't registry packages — skip.
        if arg.contains("://") || lower.starts_with("git+") {
            continue;
        }
        if arg.contains('/') || arg.contains('\\') || arg.starts_with('.') {
            continue;
        }

        // pip spec shapes we handle: foo==1.2.3, foo>=1.0, foo~=2.0, foo!=1.0,
        // foo[extra]==1.0, foo[a,b]>=1.0.
        let pkg_str = arg.as_str();

        // Strip extras: `foo[bar,baz]==1.0` -> name=`foo`, rest=`==1.0`.
        let (name_part, rest) = if let Some(bracket_pos) = pkg_str.find('[') {
            if let Some(close_pos) = pkg_str[bracket_pos..].find(']') {
                let name = &pkg_str[..bracket_pos];
                let after = &pkg_str[bracket_pos + close_pos + 1..];
                (name, after)
            } else {
                (pkg_str, "")
            }
        } else {
            let split_pos = pkg_str
                .find("==")
                .or_else(|| pkg_str.find(">="))
                .or_else(|| pkg_str.find("<="))
                .or_else(|| pkg_str.find("~="))
                .or_else(|| pkg_str.find("!="))
                .or_else(|| pkg_str.find('>'))
                .or_else(|| pkg_str.find('<'));
            if let Some(pos) = split_pos {
                (&pkg_str[..pos], &pkg_str[pos..])
            } else {
                (pkg_str, "")
            }
        };

        if name_part.is_empty() {
            continue;
        }

        let version = extract_pip_version(rest);
        let normalized = normalize_pypi_name(name_part);

        packages.push(PackageRef {
            ecosystem: Ecosystem::PyPI,
            name: normalized,
            version,
        });
    }
}

/// Normalize a PyPI package name: lowercase, replace `_` and `.` with `-`.
fn normalize_pypi_name(name: &str) -> String {
    name.to_lowercase()
        .chars()
        .map(|c| if c == '_' || c == '.' { '-' } else { c })
        .collect()
}

/// Extract exact version from pip version specifier.
/// Only returns a version for `==` (exact match).
fn extract_pip_version(spec: &str) -> Option<String> {
    if let Some(ver) = spec.strip_prefix("==") {
        let v = ver.trim();
        if !v.is_empty() {
            return Some(v.to_string());
        }
    }
    None
}

/// Flags for npm/yarn/pnpm that consume the next argument.
const NPM_ARG_FLAGS: &[&str] = &[
    "--registry",
    "--tag",
    "--scope",
    "--otp",
    "--workspace",
    "-w",
    "--prefix",
];

fn extract_npm_packages(cmd_name: &str, args: &[String], packages: &mut Vec<PackageRef>) {
    let mut iter = args.iter().peekable();
    let mut found_subcmd = false;

    // npx is special: `npx foo` runs `foo` directly. `--package`/-p can override,
    // in which case the first positional is an entry point, not a package.
    if cmd_name == "npx" {
        let mut has_explicit_package = false;
        while let Some(arg) = iter.next() {
            if arg.starts_with('-') {
                if arg == "--package" || arg == "-p" {
                    if let Some(pkg_arg) = iter.next() {
                        if let Some(pr) = parse_npm_package_spec(pkg_arg) {
                            packages.push(pr);
                            has_explicit_package = true;
                        }
                    }
                }
                continue;
            }
            if !has_explicit_package {
                if let Some(pr) = parse_npm_package_spec(arg) {
                    packages.push(pr);
                }
            }
            break;
        }
        return;
    }

    while let Some(arg) = iter.next() {
        let lower = arg.to_lowercase();
        if !found_subcmd {
            if matches!(lower.as_str(), "install" | "i" | "add") {
                found_subcmd = true;
            }
            continue;
        }

        if arg.starts_with('-') {
            let lower_ref = lower.as_str();
            if NPM_ARG_FLAGS.contains(&lower_ref) {
                let _ = iter.next();
            }
            continue;
        }

        if arg.contains("://") || arg.starts_with('.') || arg.starts_with('/') {
            continue;
        }

        if let Some(pr) = parse_npm_package_spec(arg) {
            packages.push(pr);
        }
    }
}

/// Parse an npm-style package spec: `@scope/name@version` or `name@version`.
fn parse_npm_package_spec(spec: &str) -> Option<PackageRef> {
    if spec.is_empty() {
        return None;
    }

    let (name, version) = if spec.starts_with('@') {
        // Scoped package: @scope/name@version
        // Find the version @ after the scope
        if let Some(slash_pos) = spec.find('/') {
            let after_scope = &spec[slash_pos + 1..];
            if let Some(at_pos) = after_scope.find('@') {
                let full_name = &spec[..slash_pos + 1 + at_pos];
                let ver = &after_scope[at_pos + 1..];
                (full_name, if ver.is_empty() { None } else { Some(ver) })
            } else {
                (spec, None)
            }
        } else {
            // Invalid scoped package (no slash)
            return None;
        }
    } else if let Some(at_pos) = spec.find('@') {
        let name = &spec[..at_pos];
        let ver = &spec[at_pos + 1..];
        (name, if ver.is_empty() { None } else { Some(ver) })
    } else {
        (spec, None)
    };

    if name.is_empty() {
        return None;
    }

    Some(PackageRef {
        ecosystem: Ecosystem::Npm,
        name: name.to_string(),
        version: version.map(|v| v.to_string()),
    })
}

fn extract_cargo_packages(args: &[String], packages: &mut Vec<PackageRef>) {
    let mut iter = args.iter();
    let mut found_subcmd = false;

    while let Some(arg) = iter.next() {
        let lower = arg.to_lowercase();
        if !found_subcmd {
            if matches!(lower.as_str(), "install" | "add") {
                found_subcmd = true;
            }
            continue;
        }

        if arg.starts_with('-') {
            if matches!(
                lower.as_str(),
                "--version"
                    | "--vers"
                    | "--git"
                    | "--branch"
                    | "--tag"
                    | "--rev"
                    | "--path"
                    | "--registry"
                    | "--index"
                    | "--features"
                    | "-F"
                    | "--target-dir"
                    | "--root"
                    | "--jobs"
                    | "-j"
                    | "--rename"
            ) {
                // `--version` / `--vers` carries the version of the previously seen package.
                if lower == "--version" || lower == "--vers" {
                    if let Some(ver) = iter.next() {
                        if let Some(last) = packages.last_mut() {
                            if last.ecosystem == Ecosystem::Crates && last.version.is_none() {
                                last.version = Some(ver.to_string());
                            }
                        }
                    }
                } else {
                    let _ = iter.next();
                }
                continue;
            }
            continue;
        }

        // git URLs and local paths aren't crates.io packages.
        if arg.contains("://") || arg.starts_with('.') || arg.contains('/') {
            continue;
        }

        // `cargo add foo@1.0.0` form.
        let (name, version) = split_name_version(arg, '@');

        if !name.is_empty() {
            packages.push(PackageRef {
                ecosystem: Ecosystem::Crates,
                name: name.to_string(),
                version,
            });
        }
    }
}

fn extract_gem_packages(args: &[String], packages: &mut Vec<PackageRef>) {
    let mut iter = args.iter();
    let mut found_install = false;

    while let Some(arg) = iter.next() {
        let lower = arg.to_lowercase();
        if !found_install {
            if lower == "install" {
                found_install = true;
            }
            continue;
        }

        if arg.starts_with('-') {
            if matches!(
                lower.as_str(),
                "--version" | "-v" | "--source" | "--platform" | "--install-dir" | "-i"
            ) {
                // `--version` / `-v` carries the version of the previously seen gem.
                if lower == "--version" || lower == "-v" {
                    if let Some(ver) = iter.next() {
                        if let Some(last) = packages.last_mut() {
                            if last.ecosystem == Ecosystem::RubyGems && last.version.is_none() {
                                last.version = Some(ver.to_string());
                            }
                        }
                    }
                } else {
                    let _ = iter.next();
                }
                continue;
            }
            continue;
        }

        // `gem install rails:7.0` form (also accepts bare name).
        let (name, version) = split_name_version(arg, ':');

        if !name.is_empty() {
            packages.push(PackageRef {
                ecosystem: Ecosystem::RubyGems,
                name: name.to_string(),
                version,
            });
        }
    }
}

fn extract_go_packages(args: &[String], packages: &mut Vec<PackageRef>) {
    let mut found_subcmd = false;

    for arg in args {
        let lower = arg.to_lowercase();
        if !found_subcmd {
            if matches!(lower.as_str(), "get" | "install") {
                found_subcmd = true;
            }
            continue;
        }

        if arg.starts_with('-') {
            continue;
        }

        // `go get github.com/user/pkg@v1.2.3` form.
        let (name, version) = split_name_version(arg, '@');

        if !name.is_empty() {
            packages.push(PackageRef {
                ecosystem: Ecosystem::Go,
                name: name.to_string(),
                version,
            });
        }
    }
}

fn extract_composer_packages(args: &[String], packages: &mut Vec<PackageRef>) {
    let mut found_require = false;

    for arg in args {
        if !found_require {
            if arg.to_lowercase() == "require" {
                found_require = true;
            }
            continue;
        }

        if arg.starts_with('-') {
            continue;
        }

        // `composer require vendor/package:^1.0` form.
        let (name, version) = split_name_version(arg, ':');

        if !name.is_empty() {
            packages.push(PackageRef {
                ecosystem: Ecosystem::Packagist,
                name: name.to_string(),
                version,
            });
        }
    }
}

fn extract_dotnet_packages(args: &[String], packages: &mut Vec<PackageRef>) {
    let mut iter = args.iter();
    let mut found_add = false;
    let mut found_package = false;

    while let Some(arg) = iter.next() {
        let lower = arg.to_lowercase();
        if !found_add {
            if lower == "add" {
                found_add = true;
            }
            continue;
        }

        // `dotnet add package <name>` — skip the project file arg
        if !found_package {
            if lower == "package" {
                found_package = true;
            }
            continue;
        }

        if arg.starts_with('-') {
            if lower == "--version" || lower == "-v" {
                if let Some(ver) = iter.next() {
                    if let Some(last) = packages.last_mut() {
                        if last.ecosystem == Ecosystem::NuGet && last.version.is_none() {
                            last.version = Some(ver.to_string());
                        }
                    }
                }
                continue;
            }
            if matches!(lower.as_str(), "--source" | "-s" | "--framework" | "-f") {
                let _ = iter.next();
            }
            continue;
        }

        packages.push(PackageRef {
            ecosystem: Ecosystem::NuGet,
            name: arg.to_string(),
            version: None,
        });
    }
}

/// Extract Maven/Gradle dependency coordinates from command arguments.
///
/// Handles `mvn dependency:get -Dartifact=group:artifact:version` and
/// `gradle` dependency notation `group:artifact:version`.
fn extract_maven_packages(args: &[String], packages: &mut Vec<PackageRef>) {
    for arg in args {
        // mvn form: `-Dartifact=group:artifact:version[:packaging[:classifier]]`.
        if let Some(coord) = arg.strip_prefix("-Dartifact=") {
            let parts: Vec<&str> = coord.splitn(4, ':').collect();
            if parts.len() >= 2 {
                let name = format!("{}:{}", parts[0], parts[1]);
                let version = parts.get(2).and_then(|v| {
                    if v.is_empty() {
                        None
                    } else {
                        Some(v.to_string())
                    }
                });
                packages.push(PackageRef {
                    ecosystem: Ecosystem::Maven,
                    name,
                    version,
                });
            }
            continue;
        }

        if arg.starts_with('-') {
            continue;
        }

        // Gradle dependency notation: `group:artifact:version` (at least one colon required).
        let parts: Vec<&str> = arg.splitn(4, ':').collect();
        if parts.len() >= 2 && !parts[0].is_empty() && !parts[1].is_empty() {
            let name = format!("{}:{}", parts[0], parts[1]);
            let version = parts.get(2).and_then(|v| {
                if v.is_empty() {
                    None
                } else {
                    Some(v.to_string())
                }
            });
            packages.push(PackageRef {
                ecosystem: Ecosystem::Maven,
                name,
                version,
            });
        }
    }
}

/// Extract IPv4 addresses from a shell token.
///
/// Handles:
/// - Bare IP: `1.2.3.4`
/// - user@IP: `user@1.2.3.4`
/// - IP:port: `1.2.3.4:22`
/// - user@IP:port: `user@1.2.3.4:22`
///
/// Does NOT match:
/// - IPv6 addresses
/// - Non-IP text
/// - IPs embedded inside URLs (those are handled by URL extraction)
pub fn extract_ipv4_from_token(token: &str) -> Option<Ipv4Addr> {
    let after_at = if let Some(at_pos) = token.rfind('@') {
        &token[at_pos + 1..]
    } else {
        token
    };

    // Only strip a trailing `:NNNN` — anything else after `:` would be part of
    // an IPv6 literal or something else we shouldn't touch.
    let ip_str = if let Some(colon_pos) = after_at.rfind(':') {
        let after_colon = &after_at[colon_pos + 1..];
        if !after_colon.is_empty() && after_colon.chars().all(|c| c.is_ascii_digit()) {
            &after_at[..colon_pos]
        } else {
            after_at
        }
    } else {
        after_at
    };

    // Some formats wrap with `[...]`; strip before parsing.
    let ip_str = ip_str.trim_matches(|c| c == '[' || c == ']');

    ip_str.parse::<Ipv4Addr>().ok()
}

/// Map threat-DB confidence to finding severity.
fn confidence_to_severity(c: threatdb::Confidence) -> Severity {
    match c {
        threatdb::Confidence::Confirmed => Severity::Critical,
        threatdb::Confidence::Medium => Severity::Medium,
        threatdb::Confidence::Low => Severity::Medium,
    }
}

fn hostname_rule_for_source(source: threatdb::ThreatSource) -> (RuleId, Severity, &'static str) {
    match source {
        threatdb::ThreatSource::Urlhaus => (
            RuleId::ThreatMaliciousUrl,
            Severity::High,
            "malicious hostname",
        ),
        threatdb::ThreatSource::PhishingArmy | threatdb::ThreatSource::PhishTank => (
            RuleId::ThreatPhishingUrl,
            Severity::High,
            "phishing hostname",
        ),
        threatdb::ThreatSource::ThreatFoxIoc => {
            (RuleId::ThreatThreatFoxIoc, Severity::High, "IOC hostname")
        }
        // Package/IP-oriented sources and FireHOL aren't expected on hostname
        // records, but enumerate them explicitly so the compiler flags any new
        // variant added later instead of falling through a `_` arm.
        threatdb::ThreatSource::OssfMalicious
        | threatdb::ThreatSource::DatadogMalicious
        | threatdb::ThreatSource::FeodoTracker
        | threatdb::ThreatSource::EcosystemsTyposquat
        | threatdb::ThreatSource::CisaKev
        | threatdb::ThreatSource::FireholIp
        | threatdb::ThreatSource::TorExit => (
            RuleId::ThreatMaliciousUrl,
            Severity::High,
            "malicious hostname",
        ),
    }
}

fn ip_rule_for_source(source: threatdb::ThreatSource) -> (RuleId, Severity, &'static str) {
    match source {
        threatdb::ThreatSource::TorExit => {
            (RuleId::ThreatTorExitNode, Severity::Medium, "Tor exit node")
        }
        threatdb::ThreatSource::ThreatFoxIoc => {
            (RuleId::ThreatThreatFoxIoc, Severity::High, "IOC IP")
        }
        // Same exhaustive-match rationale as the hostname mapping above.
        threatdb::ThreatSource::OssfMalicious
        | threatdb::ThreatSource::DatadogMalicious
        | threatdb::ThreatSource::FeodoTracker
        | threatdb::ThreatSource::EcosystemsTyposquat
        | threatdb::ThreatSource::CisaKev
        | threatdb::ThreatSource::Urlhaus
        | threatdb::ThreatSource::PhishingArmy
        | threatdb::ThreatSource::PhishTank
        | threatdb::ThreatSource::FireholIp => {
            (RuleId::ThreatMaliciousIp, Severity::High, "malicious IP")
        }
    }
}

/// Check input against the local threat intelligence database.
///
/// Fail-open: if `db` is `None` (no DB file loaded), returns an empty Vec
/// and does not block the command. All lookups are in-memory binary search
/// with no network I/O.
pub fn check(
    input: &str,
    shell: ShellType,
    extracted: &[ExtractedUrl],
    db: Option<&ThreatDb>,
) -> Vec<Finding> {
    let db = match db {
        Some(d) => d,
        // Fail-open: no DB loaded means no findings, never block the user.
        None => return Vec::new(),
    };

    let mut findings = Vec::new();

    let segments = crate::tokenize::tokenize(input, shell);
    let packages = extract_packages(&segments);

    for pkg in &packages {
        let db_eco = pkg.ecosystem;

        if let Some(m) = db.check_package(db_eco, &pkg.name, pkg.version.as_deref()) {
            findings.push(Finding {
                rule_id: RuleId::ThreatMaliciousPackage,
                severity: confidence_to_severity(m.confidence),
                title: format!("Known malicious {} package: {}", pkg.ecosystem, pkg.name),
                description: format!(
                    "Package '{}' in {} is flagged as malicious by {}. {}",
                    pkg.name,
                    pkg.ecosystem,
                    m.source.label(),
                    if m.all_versions_malicious {
                        "All versions are affected."
                    } else {
                        "Specific version(s) affected."
                    }
                ),
                evidence: vec![Evidence::ThreatIntel {
                    source: m.source.label().to_string(),
                    threat_type: "malicious_package".to_string(),
                    confidence: m.confidence,
                    reference: m.reference_url,
                }],
                human_view: None,
                agent_view: None,
                mitre_id: None,
                custom_rule_id: None,
            });
            // A confirmed-malicious package already explains itself — don't pile on typosquat findings.
            continue;
        }

        if let Some(t) = db.check_typosquat(db_eco, &pkg.name) {
            findings.push(Finding {
                rule_id: RuleId::ThreatPackageTyposquat,
                severity: Severity::High,
                title: format!("Confirmed typosquat: {} → {}", pkg.name, t.target_name),
                description: format!(
                    "Package '{}' in {} is a confirmed typosquat of '{}' \
                     (source: ecosyste.ms typosquatting dataset).",
                    pkg.name, pkg.ecosystem, t.target_name
                ),
                evidence: vec![Evidence::ThreatIntel {
                    source: "ecosyste.ms Typosquats".to_string(),
                    threat_type: "typosquat".to_string(),
                    confidence: threatdb::Confidence::Confirmed,
                    reference: None,
                }],
                human_view: None,
                agent_view: None,
                mitre_id: None,
                custom_rule_id: None,
            });
        }

        if let Some((popular_name, distance)) = db.check_popular_distance(db_eco, &pkg.name) {
            findings.push(Finding {
                rule_id: RuleId::ThreatPackageSimilarName,
                severity: Severity::Medium,
                title: format!(
                    "Package name similar to popular package: {} ≈ {}",
                    pkg.name, popular_name
                ),
                description: format!(
                    "Package '{}' in {} is within edit distance {} of popular package '{}'. \
                     This could indicate a typosquatting attempt.",
                    pkg.name, pkg.ecosystem, distance, popular_name
                ),
                evidence: vec![Evidence::ThreatIntel {
                    source: "popular package names".to_string(),
                    threat_type: "similar_name".to_string(),
                    confidence: threatdb::Confidence::Low,
                    reference: None,
                }],
                human_view: None,
                agent_view: None,
                mitre_id: None,
                custom_rule_id: None,
            });
        }
    }

    let mut checked_ips = std::collections::HashSet::new();
    for url_info in extracted {
        if let Some(host) = url_info.parsed.host() {
            if let Some(m) = db.check_hostname(host) {
                let (rule_id, severity, threat_type) = hostname_rule_for_source(m.source);
                findings.push(Finding {
                    rule_id,
                    severity,
                    title: format!("Threat intelligence hostname match: {}", host),
                    description: format!(
                        "Hostname '{}' appears in threat intelligence feed ({}).",
                        host,
                        m.source.label()
                    ),
                    evidence: vec![Evidence::ThreatIntel {
                        source: m.source.label().to_string(),
                        threat_type: threat_type.to_string(),
                        confidence: m.confidence,
                        reference: m.reference_url,
                    }],
                    human_view: None,
                    agent_view: None,
                    mitre_id: None,
                    custom_rule_id: None,
                });
            }

            // URL host may itself be an IP literal (e.g. `https://203.0.113.50/payload`).
            if let Ok(ip) = host.parse::<std::net::Ipv4Addr>() {
                if checked_ips.insert(ip) {
                    if let Some(m) = db.check_ip(ip) {
                        let (rule_id, severity, threat_type) = ip_rule_for_source(m.source);
                        findings.push(Finding {
                            rule_id,
                            severity,
                            title: format!("Threat intelligence IP match in URL: {}", ip),
                            description: format!(
                                "IP address {} (from URL) is flagged by {}.",
                                ip,
                                m.source.label()
                            ),
                            evidence: vec![Evidence::ThreatIntel {
                                source: m.source.label().to_string(),
                                threat_type: threat_type.to_string(),
                                confidence: m.confidence,
                                reference: m.reference_url,
                            }],
                            human_view: None,
                            agent_view: None,
                            mitre_id: None,
                            custom_rule_id: None,
                        });
                    }
                }
            }
        }
    }

    // IP literals in command tokens — ssh/scp/nc and friends.
    for seg in &segments {
        for arg in &seg.args {
            if let Some(ip) = extract_ipv4_from_token(arg) {
                if checked_ips.insert(ip) {
                    if let Some(m) = db.check_ip(ip) {
                        let (rule_id, severity, threat_type) = ip_rule_for_source(m.source);
                        findings.push(Finding {
                            rule_id,
                            severity,
                            title: format!("Threat intelligence IP match: {}", ip),
                            description: format!(
                                "IP address {} is flagged by {}.",
                                ip,
                                m.source.label()
                            ),
                            evidence: vec![Evidence::ThreatIntel {
                                source: m.source.label().to_string(),
                                threat_type: threat_type.to_string(),
                                confidence: m.confidence,
                                reference: m.reference_url,
                            }],
                            human_view: None,
                            agent_view: None,
                            mitre_id: None,
                            custom_rule_id: None,
                        });
                    }
                }
            }
        }
    }

    findings
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tokenize;

    fn tokenize_and_extract(input: &str) -> Vec<PackageRef> {
        let segments = tokenize::tokenize(input, ShellType::Posix);
        extract_packages(&segments)
    }

    #[test]
    fn pip_install_single() {
        let pkgs = tokenize_and_extract("pip install requests");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::PyPI);
        assert_eq!(pkgs[0].name, "requests");
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn pip_install_with_version() {
        let pkgs = tokenize_and_extract("pip install requests==2.31.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "requests");
        assert_eq!(pkgs[0].version, Some("2.31.0".to_string()));
    }

    #[test]
    fn pip_install_version_range_not_exact() {
        let pkgs = tokenize_and_extract("pip install requests>=2.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "requests");
        // Only `==` pins an exact version; ranges leave version as None.
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn pip3_install() {
        let pkgs = tokenize_and_extract("pip3 install flask");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::PyPI);
        assert_eq!(pkgs[0].name, "flask");
    }

    #[test]
    fn uv_install() {
        let pkgs = tokenize_and_extract("uv install numpy");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::PyPI);
        assert_eq!(pkgs[0].name, "numpy");
    }

    #[test]
    fn pip_install_multiple() {
        let pkgs = tokenize_and_extract("pip install requests flask django");
        assert_eq!(pkgs.len(), 3);
        assert_eq!(pkgs[0].name, "requests");
        assert_eq!(pkgs[1].name, "flask");
        assert_eq!(pkgs[2].name, "django");
    }

    #[test]
    fn pip_install_with_extras() {
        let pkgs = tokenize_and_extract("pip install requests[security]==2.31.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "requests");
        assert_eq!(pkgs[0].version, Some("2.31.0".to_string()));
    }

    #[test]
    fn pip_install_skips_flags() {
        let pkgs =
            tokenize_and_extract("pip install --index-url https://pypi.org/simple/ requests");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "requests");
    }

    #[test]
    fn pip_install_skips_url() {
        let pkgs =
            tokenize_and_extract("pip install git+https://github.com/user/repo.git requests");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "requests");
    }

    #[test]
    fn pip_install_skips_local_path() {
        let pkgs = tokenize_and_extract("pip install ./local_pkg requests");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "requests");
    }

    #[test]
    fn pip_normalizes_name() {
        let pkgs = tokenize_and_extract("pip install My_Package.Name");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "my-package-name");
    }

    #[test]
    fn pip_no_install_subcommand() {
        let pkgs = tokenize_and_extract("pip freeze");
        assert!(pkgs.is_empty());
    }

    #[test]
    fn npm_install_single() {
        let pkgs = tokenize_and_extract("npm install lodash");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Npm);
        assert_eq!(pkgs[0].name, "lodash");
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn npm_install_with_version() {
        let pkgs = tokenize_and_extract("npm install lodash@4.17.21");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "lodash");
        assert_eq!(pkgs[0].version, Some("4.17.21".to_string()));
    }

    #[test]
    fn npm_install_scoped() {
        let pkgs = tokenize_and_extract("npm install @angular/core@16.0.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "@angular/core");
        assert_eq!(pkgs[0].version, Some("16.0.0".to_string()));
    }

    #[test]
    fn npm_install_scoped_no_version() {
        let pkgs = tokenize_and_extract("npm install @types/node");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "@types/node");
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn npm_i_shorthand() {
        let pkgs = tokenize_and_extract("npm i express");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "express");
    }

    #[test]
    fn yarn_add() {
        let pkgs = tokenize_and_extract("yarn add react@18.2.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Npm);
        assert_eq!(pkgs[0].name, "react");
        assert_eq!(pkgs[0].version, Some("18.2.0".to_string()));
    }

    #[test]
    fn pnpm_add() {
        let pkgs = tokenize_and_extract("pnpm add vue");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Npm);
        assert_eq!(pkgs[0].name, "vue");
    }

    #[test]
    fn bun_add() {
        let pkgs = tokenize_and_extract("bun add elysia");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Npm);
        assert_eq!(pkgs[0].name, "elysia");
    }

    #[test]
    fn npx_package() {
        let pkgs = tokenize_and_extract("npx create-react-app my-app");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Npm);
        assert_eq!(pkgs[0].name, "create-react-app");
    }

    #[test]
    fn npx_scoped_package() {
        let pkgs = tokenize_and_extract("npx @angular/cli new my-app");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "@angular/cli");
    }

    #[test]
    fn npx_with_package_flag() {
        let pkgs = tokenize_and_extract("npx --package typescript tsc");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "typescript");
    }

    #[test]
    fn npm_install_skips_save_dev() {
        let pkgs = tokenize_and_extract("npm install --save-dev jest");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "jest");
    }

    #[test]
    fn npm_install_multiple() {
        let pkgs = tokenize_and_extract("npm install react react-dom");
        assert_eq!(pkgs.len(), 2);
        assert_eq!(pkgs[0].name, "react");
        assert_eq!(pkgs[1].name, "react-dom");
    }

    #[test]
    fn cargo_install() {
        let pkgs = tokenize_and_extract("cargo install ripgrep");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Crates);
        assert_eq!(pkgs[0].name, "ripgrep");
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn cargo_add() {
        let pkgs = tokenize_and_extract("cargo add serde");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "serde");
    }

    #[test]
    fn cargo_add_with_version() {
        let pkgs = tokenize_and_extract("cargo add serde@1.0.193");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "serde");
        assert_eq!(pkgs[0].version, Some("1.0.193".to_string()));
    }

    #[test]
    fn cargo_install_with_version_flag() {
        let pkgs = tokenize_and_extract("cargo install ripgrep --version 14.0.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "ripgrep");
        assert_eq!(pkgs[0].version, Some("14.0.0".to_string()));
    }

    #[test]
    fn cargo_skips_git_url() {
        let pkgs = tokenize_and_extract("cargo install --git https://github.com/user/repo");
        assert!(pkgs.is_empty());
    }

    #[test]
    fn cargo_build_not_install() {
        let pkgs = tokenize_and_extract("cargo build --release");
        assert!(pkgs.is_empty());
    }

    #[test]
    fn gem_install() {
        let pkgs = tokenize_and_extract("gem install rails");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::RubyGems);
        assert_eq!(pkgs[0].name, "rails");
    }

    #[test]
    fn gem_install_with_version_flag() {
        let pkgs = tokenize_and_extract("gem install rails --version 7.0.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "rails");
        assert_eq!(pkgs[0].version, Some("7.0.0".to_string()));
    }

    #[test]
    fn gem_install_with_colon_version() {
        let pkgs = tokenize_and_extract("gem install rails:7.0.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "rails");
        assert_eq!(pkgs[0].version, Some("7.0.0".to_string()));
    }

    #[test]
    fn go_get() {
        let pkgs = tokenize_and_extract("go get github.com/gin-gonic/gin");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Go);
        assert_eq!(pkgs[0].name, "github.com/gin-gonic/gin");
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn go_get_with_version() {
        let pkgs = tokenize_and_extract("go get github.com/gin-gonic/gin@v1.9.1");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "github.com/gin-gonic/gin");
        assert_eq!(pkgs[0].version, Some("v1.9.1".to_string()));
    }

    #[test]
    fn go_install() {
        let pkgs = tokenize_and_extract("go install golang.org/x/tools/gopls@latest");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "golang.org/x/tools/gopls");
        assert_eq!(pkgs[0].version, Some("latest".to_string()));
    }

    #[test]
    fn composer_require() {
        let pkgs = tokenize_and_extract("composer require monolog/monolog");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::Packagist);
        assert_eq!(pkgs[0].name, "monolog/monolog");
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn composer_require_with_version() {
        let pkgs = tokenize_and_extract("composer require monolog/monolog:^3.0");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "monolog/monolog");
        assert_eq!(pkgs[0].version, Some("^3.0".to_string()));
    }

    #[test]
    fn dotnet_add_package() {
        let pkgs = tokenize_and_extract("dotnet add package Newtonsoft.Json");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::NuGet);
        assert_eq!(pkgs[0].name, "Newtonsoft.Json");
        assert_eq!(pkgs[0].version, None);
    }

    #[test]
    fn dotnet_add_package_with_version() {
        let pkgs = tokenize_and_extract("dotnet add package Newtonsoft.Json --version 13.0.3");
        assert_eq!(pkgs.len(), 1);
        assert_eq!(pkgs[0].name, "Newtonsoft.Json");
        assert_eq!(pkgs[0].version, Some("13.0.3".to_string()));
    }

    #[test]
    fn no_packages_in_ls() {
        let pkgs = tokenize_and_extract("ls -la");
        assert!(pkgs.is_empty());
    }

    #[test]
    fn no_packages_in_echo() {
        let pkgs = tokenize_and_extract("echo hello world");
        assert!(pkgs.is_empty());
    }

    #[test]
    fn piped_commands_both_extracted() {
        let pkgs = tokenize_and_extract("pip install requests && npm install lodash");
        assert_eq!(pkgs.len(), 2);
        assert_eq!(pkgs[0].ecosystem, Ecosystem::PyPI);
        assert_eq!(pkgs[0].name, "requests");
        assert_eq!(pkgs[1].ecosystem, Ecosystem::Npm);
        assert_eq!(pkgs[1].name, "lodash");
    }

    #[test]
    fn ipv4_bare() {
        let ip = extract_ipv4_from_token("1.2.3.4");
        assert_eq!(ip, Some(Ipv4Addr::new(1, 2, 3, 4)));
    }

    #[test]
    fn ipv4_with_user() {
        let ip = extract_ipv4_from_token("user@192.168.1.1");
        assert_eq!(ip, Some(Ipv4Addr::new(192, 168, 1, 1)));
    }

    #[test]
    fn ipv4_with_port() {
        let ip = extract_ipv4_from_token("10.0.0.1:22");
        assert_eq!(ip, Some(Ipv4Addr::new(10, 0, 0, 1)));
    }

    #[test]
    fn ipv4_with_user_and_port() {
        let ip = extract_ipv4_from_token("root@10.0.0.1:22");
        assert_eq!(ip, Some(Ipv4Addr::new(10, 0, 0, 1)));
    }

    #[test]
    fn ipv4_localhost() {
        let ip = extract_ipv4_from_token("127.0.0.1");
        assert_eq!(ip, Some(Ipv4Addr::new(127, 0, 0, 1)));
    }

    #[test]
    fn ipv4_not_an_ip() {
        let ip = extract_ipv4_from_token("hello");
        assert!(ip.is_none());
    }

    #[test]
    fn ipv4_partial_not_valid() {
        let ip = extract_ipv4_from_token("1.2.3");
        assert!(ip.is_none());
    }

    #[test]
    fn ipv4_out_of_range() {
        let ip = extract_ipv4_from_token("999.999.999.999");
        assert!(ip.is_none());
    }

    #[test]
    fn ipv6_not_matched() {
        assert!(extract_ipv4_from_token("::1").is_none());
        assert!(extract_ipv4_from_token("2001:db8::1").is_none());
        assert!(extract_ipv4_from_token("fe80::1%eth0").is_none());
    }

    #[test]
    fn ipv4_empty_string() {
        let ip = extract_ipv4_from_token("");
        assert!(ip.is_none());
    }

    #[test]
    fn ipv4_in_brackets() {
        let ip = extract_ipv4_from_token("[10.0.0.1]");
        assert_eq!(ip, Some(Ipv4Addr::new(10, 0, 0, 1)));
    }

    #[test]
    fn check_returns_empty_without_db() {
        let findings = check("pip install malicious-pkg", ShellType::Posix, &[], None);
        assert!(findings.is_empty(), "check() must be fail-open without DB");
    }

    #[test]
    fn hostname_rule_urlhaus_maps_to_malicious_url() {
        let (rule, sev, _) = hostname_rule_for_source(threatdb::ThreatSource::Urlhaus);
        assert_eq!(rule, RuleId::ThreatMaliciousUrl);
        assert_eq!(sev, Severity::High);
    }

    #[test]
    fn hostname_rule_phishing_sources_map_to_phishing_url() {
        for source in [
            threatdb::ThreatSource::PhishingArmy,
            threatdb::ThreatSource::PhishTank,
        ] {
            let (rule, sev, _) = hostname_rule_for_source(source);
            assert_eq!(rule, RuleId::ThreatPhishingUrl);
            assert_eq!(sev, Severity::High);
        }
    }

    #[test]
    fn hostname_rule_threatfox_maps_to_ioc() {
        let (rule, sev, _) = hostname_rule_for_source(threatdb::ThreatSource::ThreatFoxIoc);
        assert_eq!(rule, RuleId::ThreatThreatFoxIoc);
        assert_eq!(sev, Severity::High);
    }

    #[test]
    fn ip_rule_tor_exit_maps_to_medium() {
        let (rule, sev, _) = ip_rule_for_source(threatdb::ThreatSource::TorExit);
        assert_eq!(rule, RuleId::ThreatTorExitNode);
        assert_eq!(sev, Severity::Medium);
    }

    #[test]
    fn ip_rule_threatfox_maps_to_ioc() {
        let (rule, sev, _) = ip_rule_for_source(threatdb::ThreatSource::ThreatFoxIoc);
        assert_eq!(rule, RuleId::ThreatThreatFoxIoc);
        assert_eq!(sev, Severity::High);
    }

    #[test]
    fn ip_rule_feodo_maps_to_malicious_ip() {
        let (rule, sev, _) = ip_rule_for_source(threatdb::ThreatSource::FeodoTracker);
        assert_eq!(rule, RuleId::ThreatMaliciousIp);
        assert_eq!(sev, Severity::High);
    }
}