use crate::acquire::sources::RepoSource;
use crate::acquire::strategy::{AcquireOptions, AcquisitionStrategy};
use crate::acquire::zip_only::ZipOnlyStrategy;
use crate::acquire::zip_with_history::ZipWithHistoryStrategy;
use crate::cli::UI;
use crate::core::{Config, Finding, Severity};
use crate::redteam::bridge::BridgeError;
use crate::redteam::findings::{parse_mcp_scan_findings, parse_redteam_findings};
use crate::redteam::RedteamBridge;
use anyhow::Result;
use std::path::PathBuf;
pub struct AcquireParams {
pub source: String,
pub output: Option<PathBuf>,
pub strategy: String,
pub no_history: bool,
pub fail_on: String,
pub force: bool,
pub opts: AcquireOptions,
}
pub async fn execute(params: AcquireParams, ui: &UI) -> Result<()> {
ui.header("Acquisition");
ui.blank();
let repo_source = RepoSource::parse(¶ms.source)?;
ui.field("Source", &repo_source);
let target = params.output.unwrap_or_else(|| {
let name = match &repo_source {
RepoSource::GitHub { repo, .. } => repo.clone(),
RepoSource::GitLab { repo, .. } => repo.clone(),
RepoSource::Bitbucket { repo, .. } => repo.clone(),
RepoSource::DirectUrl { .. } => "download".to_string(),
};
PathBuf::from(name)
});
ui.field("Target", target.display());
ui.blank();
if target.exists() && !params.force {
anyhow::bail!("Target directory already exists. Use --force to overwrite.");
}
let config = Config::default();
let llm_security = config.llm_security.clone();
let actual_strategy = if params.no_history {
"zip-only"
} else {
¶ms.strategy
};
ui.field("Strategy", actual_strategy);
if let Some(depth) = params.opts.depth {
ui.field("Depth", depth);
}
if params.opts.recurse_submodules {
ui.field("Submodules", "recursive");
}
if params.opts.lfs {
ui.field("LFS", "enabled");
}
ui.blank();
let report = match actual_strategy {
"zip-with-history" => {
let strat = ZipWithHistoryStrategy::new(config);
strat.acquire(&repo_source, &target, ¶ms.opts).await?
}
"zip-only" => {
let strat = ZipOnlyStrategy::new(config);
strat.acquire(&repo_source, &target, ¶ms.opts).await?
}
_ => anyhow::bail!("Unsupported strategy: {}", actual_strategy),
};
let mut audit_findings: Vec<Finding> = Vec::new();
if llm_security.enabled && llm_security.scan_on_acquire {
let bridge = RedteamBridge::new(llm_security.binary.clone());
if bridge.is_available() {
ui.field("LLM Security", "scanning...");
for pattern in &llm_security.mcp_config_patterns {
let config_path = target.join(pattern);
if config_path.exists() {
let path_str = config_path.display().to_string();
match bridge
.scan_mcp_server("cat", std::slice::from_ref(&path_str))
.await
{
Ok(result) => {
audit_findings
.extend(parse_mcp_scan_findings(&result, Some(&path_str)));
}
Err(BridgeError::NotInstalled(_)) => {}
Err(e) => {
ui.warning(format!("MCP config scan failed: {}", e));
}
}
}
}
if let Some(model_id) = repo_source.hf_model_id() {
match bridge.pipeline_scan(model_id, None).await {
Ok(result) => {
audit_findings.extend(parse_redteam_findings(&result));
}
Err(BridgeError::NotInstalled(_)) => {}
Err(e) => {
ui.warning(format!("HF model audit failed: {}", e));
}
}
}
let readme = target.join("README.md");
if readme.exists() {
if let Ok(content) = std::fs::read_to_string(&readme) {
if content.len() < 50_000 {
match bridge.firewall_check(&content).await {
Ok(result) => {
audit_findings.extend(parse_redteam_findings(&result));
}
Err(BridgeError::NotInstalled(_)) => {}
Err(e) => {
ui.warning(format!("README firewall check failed: {}", e));
}
}
}
}
}
let model_extensions = [
".pkl",
".pt",
".bin",
".pth",
".gguf",
".onnx",
".safetensors",
];
let has_model_files = walkdir::WalkDir::new(&target)
.max_depth(3)
.into_iter()
.filter_map(|e| e.ok())
.any(|e| {
let name = e.path().to_string_lossy().to_lowercase();
model_extensions.iter().any(|ext| name.ends_with(ext))
});
if has_model_files {
if llm_security.run_modelscan {
let ms = crate::toolbridges::modelscan::ModelScanBridge::new();
if ms.is_available() {
match ms.scan(&target).await {
Ok(findings) => audit_findings.extend(findings),
Err(crate::toolbridges::CliError::NotInstalled(_)) => {}
Err(e) => {
ui.warning(format!("ModelScan failed: {}", e));
}
}
}
}
if llm_security.run_picklescan {
let ps = crate::toolbridges::picklescan::PickleScanBridge::new();
if ps.is_available() {
match ps.scan(&target).await {
Ok(findings) => audit_findings.extend(findings),
Err(crate::toolbridges::CliError::NotInstalled(_)) => {}
Err(e) => {
ui.warning(format!("PickleScan failed: {}", e));
}
}
}
}
if llm_security.run_fickling {
let fk = crate::toolbridges::fickling::FicklingBridge::new();
if fk.is_available() {
match fk.scan(&target).await {
Ok(findings) => audit_findings.extend(findings),
Err(crate::toolbridges::CliError::NotInstalled(_)) => {}
Err(e) => {
ui.warning(format!("Fickling failed: {}", e));
}
}
}
}
}
if let Some(max_sev) = audit_findings.iter().map(|f| f.severity).max() {
let threshold = Severity::parse_str(¶ms.fail_on).unwrap_or(Severity::High);
if max_sev >= threshold {
let blocking_count = audit_findings
.iter()
.filter(|f| f.severity >= threshold)
.count();
ui.error(format!(
"Clone blocked: {} LLM security finding(s) at or above '{}' threshold",
blocking_count, params.fail_on
));
for f in &audit_findings {
if f.severity >= threshold {
ui.field(&format!("[{}]", f.severity), &f.title);
}
}
anyhow::bail!(
"Security threshold exceeded: findings at or above {} severity",
params.fail_on
);
}
}
}
}
if !audit_findings.is_empty() {
ui.warning(format!(
"{} LLM security finding(s) post-clone",
audit_findings.len()
));
for f in &audit_findings {
ui.field(&format!("[{}]", f.severity), &f.title);
}
}
ui.result_banner(
true,
"Acquisition complete",
&[
("Target", target.display().to_string()),
(
"History",
if report.has_history { "Yes" } else { "No" }.to_string(),
),
(
"Hooks removed",
report.sanitize_report.removed_hooks.len().to_string(),
),
],
);
let repo_url = repo_source.to_url();
let branch = repo_source.ref_name();
crate::graphrag::client::trigger_indexing(&target, Some(&repo_url), branch).await;
Ok(())
}