profileme / src/extract.rs

use crate::config::Config;
use std::path::Path;
use std::process::Command;

fn run_cmd(prog: &str, args: &[&str]) -> Option<String> {
    let out = Command::new(prog).args(args).output().ok()?;
    if !out.status.success() && out.stdout.is_empty() {
        return None;
    }
    Some(String::from_utf8_lossy(&out.stdout).to_string())
}

fn read_text_file(path: &Path, max: usize) -> Option<String> {
    let bytes = std::fs::read(path).ok()?;
    if bytes.contains(&0) {
        return None;
    }
    let mut s = String::from_utf8_lossy(&bytes).to_string();
    if s.chars().count() > max {
        s = s.chars().take(max).collect();
        s.push_str("\n...[truncated]...");
    }
    Some(s)
}

fn xml_text(xml: &str) -> String {
    // extract text nodes using quick-xml
    use quick_xml::events::Event;
    use quick_xml::Reader;
    let mut reader = Reader::from_str(xml);
    reader.config_mut().trim_text(true);
    let mut buf = Vec::new();
    let mut out = String::new();
    loop {
        match reader.read_event_into(&mut buf) {
            Ok(Event::Text(t)) => {
                if let Ok(txt) = t.unescape() {
                    let txt = txt.trim();
                    if !txt.is_empty() {
                        out.push_str(txt);
                        out.push(' ');
                    }
                }
            }
            Ok(Event::Eof) => break,
            Err(_) => break,
            _ => {}
        }
        buf.clear();
    }
    out
}

fn unzip_entry(path: &Path, entry: &str) -> Option<String> {
    run_cmd("unzip", &["-p", path.to_str()?, entry])
}

fn extract_office(path: &Path, ext: &str) -> Option<String> {
    match ext {
        "docx" => {
            let xml = unzip_entry(path, "word/document.xml")?;
            Some(xml_text(&xml))
        }
        "xlsx" => {
            // shared strings
            let mut out = String::new();
            if let Some(ss) = unzip_entry(path, "xl/sharedStrings.xml") {
                out.push_str(&xml_text(&ss));
            }
            // also try sheet1
            if let Some(sh) = unzip_entry(path, "xl/worksheets/sheet1.xml") {
                out.push(' ');
                out.push_str(&xml_text(&sh));
            }
            if out.trim().is_empty() {
                None
            } else {
                Some(out)
            }
        }
        "pptx" => {
            // extract all slides via listing
            let listing = run_cmd("unzip", &["-Z1", path.to_str()?])?;
            let mut out = String::new();
            for line in listing.lines() {
                if line.starts_with("ppt/slides/slide") && line.ends_with(".xml") {
                    if let Some(xml) = unzip_entry(path, line) {
                        out.push_str(&xml_text(&xml));
                        out.push('\n');
                    }
                }
            }
            if out.trim().is_empty() {
                None
            } else {
                Some(out)
            }
        }
        "odt" | "ods" | "odp" => {
            let xml = unzip_entry(path, "content.xml")?;
            Some(xml_text(&xml))
        }
        _ => None,
    }
}

fn extract_pdf(path: &Path) -> Option<String> {
    run_cmd("pdftotext", &["-layout", path.to_str()?, "-"])
}

fn extract_sqlite(path: &Path) -> Option<String> {
    // tables
    let tables_out = run_cmd("sqlite3", &[path.to_str()?, ".tables"])?;
    let mut out = String::new();
    out.push_str("TABLES: ");
    out.push_str(tables_out.trim());
    out.push('\n');
    for table in tables_out.split_whitespace() {
        // sanitize table name (alphanumeric/underscore)
        if !table.chars().all(|c| c.is_alphanumeric() || c == '_') {
            continue;
        }
        let q = format!("SELECT * FROM {} LIMIT 30;", table);
        if let Some(rows) = run_cmd("sqlite3", &[path.to_str()?, &q]) {
            out.push_str(&format!("\n-- {} --\n", table));
            out.push_str(&rows);
        }
    }
    Some(out)
}

fn extract_archive(path: &Path) -> Option<String> {
    // list contents only
    if let Some(s) = run_cmd("unzip", &["-l", path.to_str()?]) {
        return Some(s);
    }
    if let Some(s) = run_cmd("7z", &["l", path.to_str()?]) {
        return Some(s);
    }
    None
}

fn extract_image(path: &Path, ocr: bool) -> Option<String> {
    if !ocr {
        return None;
    }
    run_cmd("tesseract", &[path.to_str()?, "stdout"])
}

pub fn extract(path: &Path, cfg: &Config) -> Option<String> {
    let ext = path
        .extension()
        .and_then(|e| e.to_str())
        .unwrap_or("")
        .to_lowercase();

    let res = match ext.as_str() {
        "pdf" => extract_pdf(path),
        "docx" | "xlsx" | "pptx" | "odt" | "ods" | "odp" => extract_office(path, &ext),
        "db" | "sqlite" | "sqlite3" => extract_sqlite(path),
        "zip" | "tar" | "gz" | "7z" | "rar" => extract_archive(path),
        "png" | "jpg" | "jpeg" | "tif" | "tiff" | "bmp" | "webp" => extract_image(path, cfg.ocr),
        _ => {
            // try sqlite if file smells like sqlite
            if let Ok(bytes) = std::fs::read(path) {
                if bytes.starts_with(b"SQLite format 3") {
                    return extract_sqlite(path);
                }
            }
            read_text_file(path, 200_000)
        }
    };

    // truncate
    res.map(|mut s| {
        if s.chars().count() > 200_000 {
            s = s.chars().take(200_000).collect();
            s.push_str("\n...[truncated]...");
        }
        s
    })
}