QRGen/cli/src/main.rs

use anyhow::{bail, Context, Result};
use clap::{CommandFactory, Parser, Subcommand};
use clap_complete::{generate, Shell};
use indicatif::{ProgressBar, ProgressStyle};
use qr_core::qr::{QrCode, QrConfig, VersionMode};
use qr_core::text_builder;
use qr_core::version::EcLevel;
use qr_core::QrError;
use serde::Deserialize;
use std::io::{self, Read};
use std::path::Path;
use std::process;

// ──────────────────── 结构定义 ────────────────────

#[derive(Parser)]
#[command(
    name = "qrgen",
    version,
    about = "QR 码生成/解码工具 — 从零手搓的 ISO/IEC 18004 实现",
    after_help = "示例:\n  qrgen encode \"Hello\" -o qr.png\n  qrgen encode --mode wifi --ssid MyWiFi --password pass123\n  qrgen decode qr.png\n  echo \"Hello\" | qrgen encode -\n\n补全:\n  qrgen --generate-completions bash > /usr/share/bash-completion/completions/qrgen"
)]
struct Cli {
    #[command(subcommand)]
    command: Command,
    #[arg(long, value_name = "SHELL", value_parser = ["bash", "zsh", "fish", "powershell", "elvish"])]
    generate_completions: Option<String>,
}

#[derive(Subcommand)]
#[allow(clippy::large_enum_variant)]
enum Command {
    /// 编码：文本 → QR 码
    Encode {
        /// 要编码的内容（传 `-` 从 stdin 读取）
        #[arg(default_value = "-")]
        content: String,
        /// 输出文件；不指定则终端 ASCII
        #[arg(short = 'o', long)]
        output: Option<String>,
        #[command(flatten)]
        opts: EncodeOpts,
    },
    /// 解码：QR 码图片 → 文本
    Decode {
        /// 图片文件路径（传 `-` 从 stdin 读取）
        #[arg(default_value = "-")]
        file: String,
    },
}

#[derive(clap::Args, Clone)]
struct EncodeOpts {
    #[arg(short = 'l', long, default_value = "M")]
    level: String,
    #[arg(short = 'V', long, value_parser = clap::value_parser!(u8).range(1..=40))]
    version: Option<u8>,
    #[arg(short = 's', long, default_value = "4", value_parser = clap::value_parser!(u8).range(1..=20))]
    size: u8,
    #[arg(short = 'm', long, default_value = "4", value_parser = clap::value_parser!(u8).range(0..=100))]
    margin: u8,
    #[arg(long)] fg: Option<String>,
    #[arg(long)] bg: Option<String>,
    #[arg(long)] logo: Option<String>,
    #[arg(short = 'f', long, default_value = "png")]
    format: String,
    #[arg(long)] mode: Option<String>,
    // WiFi — 密码优先从环境变量 QRGEN_WIFI_PASSWORD 读取，避免命令行泄露
    #[arg(long)] ssid: Option<String>,
    #[arg(long)] password: Option<String>,
    #[arg(long, default_value = "WPA")]
    encryption: String,
    #[arg(long)] hidden: bool,
    // vCard
    #[arg(long)] name: Option<String>,
    #[arg(long)] phone: Option<String>,
    #[arg(long)] email: Option<String>,
    #[arg(long)] company: Option<String>,
    #[arg(long)] title: Option<String>,
    #[arg(long)] address: Option<String>,
    #[arg(long = "vcard-url")] vcard_url: Option<String>,
    #[arg(long)] birthday: Option<String>,
    #[arg(long)] note: Option<String>,
    #[arg(long)] photo: Option<String>,
    // Email
    #[arg(long)] to: Option<String>,
    #[arg(long)] subject: Option<String>,
    #[arg(long)] body: Option<String>,
    // Phone/SMS
    #[arg(long)] number: Option<String>,
    #[arg(long)] message: Option<String>,
    // URL
    #[arg(long)] url: Option<String>,
    // Batch
    #[arg(long)] batch: Option<String>,
    #[arg(long)] output_dir: Option<String>,
}

#[derive(Deserialize)]
struct BatchEntry {
    content: Option<String>,
    level: Option<String>,
    ssid: Option<String>,
    password: Option<String>,
    encryption: Option<String>,
    #[serde(default)]
    hidden: Option<bool>,
    name: Option<String>,
    phone: Option<String>,
    email: Option<String>,
    company: Option<String>,
    address: Option<String>,
    to: Option<String>,
    subject: Option<String>,
    body: Option<String>,
    number: Option<String>,
    message: Option<String>,
    url: Option<String>,
}

// ──────────────────── 入口 ────────────────────

fn main() {
    let cli = Cli::parse();
    if let Some(s) = cli.generate_completions {
        if let Ok(sh) = s.parse::<Shell>() {
            generate(sh, &mut Cli::command(), "qrgen", &mut io::stdout());
            return;
        }
    }
    let r = match cli.command {
        Command::Encode { content, output, opts } => cmd_encode(&content, &output, &opts),
        Command::Decode { file } => cmd_decode(&file),
    };
    if let Err(e) = r {
        eprintln!("qrgen: {:#}", e);
        process::exit(1);
    }
}

// ──────────────────── I/O 辅助 ────────────────────

/// 最大 stdin 读取量：10 MB
const STDIN_MAX: u64 = 10 * 1024 * 1024;

fn stdin_bytes() -> Result<Vec<u8>> {
    let mut b = Vec::new();
    io::stdin()
        .take(STDIN_MAX)
        .read_to_end(&mut b)
        .with_context(|| "无法读取 stdin")?;
    Ok(b)
}
fn stdin_text() -> Result<String> {
    let mut s = String::new();
    io::stdin()
        .take(STDIN_MAX)
        .read_to_string(&mut s)
        .with_context(|| "无法读取 stdin")?;
    let t = s.trim().to_string();
    if t.is_empty() {
        bail!("stdin 为空");
    }
    Ok(t)
}

// ──────────────────── 编码 ────────────────────

fn cmd_encode(content: &str, output: &Option<String>, opts: &EncodeOpts) -> Result<()> {
    let text = if content == "-" {
        stdin_text()?
    } else {
        content.to_string()
    };
    let final_text = if let Some(m) = &opts.mode {
        build_mode(m, opts, &text)?
    } else if let Some(ref bf) = opts.batch {
        check_path(bf)?; // 批量文件路径检查
        if let Some(od) = &opts.output_dir {
            check_path(od)?; // 输出目录路径检查
        }
        return do_batch(bf, opts);
    } else {
        text
    };

    let level: EcLevel = opts.level.parse().map_err(|e: QrError| anyhow::anyhow!(e))?;
    // --logo 文件路径也需安全检查
    if let Some(logo_path) = &opts.logo {
        check_path(logo_path)?;
    }
    let logo = opts
        .logo
        .as_ref()
        .map(std::fs::read)
        .transpose()
        .with_context(|| "无法读取 logo")?;

    let config = QrConfig {
        level,
        version: opts
            .version
            .map(VersionMode::Fixed)
            .unwrap_or(VersionMode::Auto),
        margin: opts.margin,
    };
    let qr = QrCode::encode(&final_text, config).map_err(|e| anyhow::anyhow!("编码失败: {e}"))?;

    let fg = opts.fg.as_deref();
    let bg = opts.bg.as_deref();

    match output {
        Some(p) => {
            check_path(p)?;
            let ext = Path::new(p)
                .extension()
                .and_then(|e| e.to_str())
                .unwrap_or("")
                .to_lowercase();
            if ext.is_empty() {
                eprintln!("警告: 无法从扩展名推断输出格式，回退到 PNG");
            }
            match ext.as_str() {
                "svg" => {
                    std::fs::write(p, qr.to_svg(logo.as_deref(), fg, bg))?;
                    eprintln!("已生成: {p} (版本 {}, SVG)", qr.version.0);
                }
                _ => {
                    let fmt = qr_core::render::image::OutputFormat::from_ext(&ext)
                        .or_else(|| qr_core::render::image::OutputFormat::from_ext(&opts.format))
                        .unwrap_or_else(|| {
                            eprintln!("警告: 无法识别扩展名 '{}'，回退到 PNG", ext);
                            qr_core::render::image::OutputFormat::Png
                        });
                    std::fs::write(
                        p,
                        qr.to_image_bytes(opts.size, logo.as_deref(), Some(fmt), fg, bg)?,
                    )?;
                    eprintln!(
                        "已生成: {p} (版本 {}, {}×{}, {:?}, {})",
                        qr.version.0,
                        qr.size(),
                        qr.size(),
                        qr.level,
                        fmt.extension()
                    );
                }
            }
        }
        None => {
            println!("{}", qr.to_ascii(false));
        }
    }
    Ok(())
}

fn build_mode(mode: &str, opts: &EncodeOpts, fb: &str) -> Result<String> {
    match mode {
        "wifi" => {
            let s = opts
                .ssid
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("WiFi 模式需要 --ssid"))?;
            // 密码优先从 --password 读取，未提供时尝试环境变量 QRGEN_WIFI_PASSWORD
            let env_pwd = std::env::var("QRGEN_WIFI_PASSWORD").ok();
            let pwd = opts
                .password
                .as_deref()
                .or(env_pwd.as_deref())
                .unwrap_or("");
            Ok(text_builder::build_wifi_text(
                s,
                pwd,
                &opts.encryption,
                opts.hidden,
            ))
        }
        "vcard" => Ok(text_builder::build_vcard_text(
            opts.name.as_deref().unwrap_or(""),
            opts.phone.as_deref().unwrap_or(""),
            opts.email.as_deref().unwrap_or(""),
            opts.company.as_deref().unwrap_or(""),
            opts.address.as_deref().unwrap_or(""),
            opts.title.as_deref().unwrap_or(""),
            opts.vcard_url.as_deref().unwrap_or(""),
            opts.birthday.as_deref().unwrap_or(""),
            opts.note.as_deref().unwrap_or(""),
            opts.photo.as_deref().unwrap_or(""),
        )),
        "email" => {
            let t = opts
                .to
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("Email 模式需要 --to"))?;
            Ok(text_builder::build_email_text(
                t,
                opts.subject.as_deref().unwrap_or(""),
                opts.body.as_deref().unwrap_or(""),
            ))
        }
        "phone" => Ok(text_builder::build_phone_text(
            opts.number
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("需要 --number"))?,
        )),
        "sms" => Ok(text_builder::build_sms_text(
            opts.number
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("需要 --number"))?,
            opts.message.as_deref().unwrap_or(""),
        )),
        "url" => opts
            .url
            .clone()
            .ok_or_else(|| anyhow::anyhow!("URL 模式需要 --url")),
        "text" => Ok(fb.to_string()),
        _m => bail!("未知模式: {_m}，支持 text/url/wifi/vcard/email/phone/sms/batch"),
    }
}

// ──────────────────── 解码 ────────────────────

fn cmd_decode(file: &str) -> Result<()> {
    if file != "-" {
        check_path(file)?;
    }
    let bytes = if file == "-" {
        stdin_bytes()?
    } else {
        std::fs::read(file)?
    };
    let r = qr_core::decoder::decode_image(&bytes)
        .map_err(|e| anyhow::anyhow!("解码失败: {e}"))?;
    println!("{}", r.text);
    eprintln!(
        "版本: {}  级别: {:?}  掩码: {}  纠错: {} 码字",
        r.version, r.level, r.mask, r.errors_corrected
    );
    Ok(())
}

// ──────────────────── 批量 ────────────────────

fn do_batch(file: &str, opts: &EncodeOpts) -> Result<()> {
    let input = std::fs::read_to_string(file)
        .with_context(|| format!("无法读取批量文件 '{file}'"))?;
    let entries: Vec<BatchEntry> = serde_json::from_str(&input)
        .or_else(|_| parse_csv(&input))
        .map_err(|e| anyhow::anyhow!("解析失败: {e}"))?;
    let out = opts.output_dir.as_deref().unwrap_or("batch_output");
    std::fs::create_dir_all(out)
        .with_context(|| format!("无法创建目录 '{out}'"))?;

    let total = entries.len();
    let pb = ProgressBar::new(total as u64);
    pb.set_style(
        ProgressStyle::default_bar()
            .template("{spinner:.green} [{bar:30.cyan/blue}] {pos}/{len} {msg}")
            .unwrap(),
    );

    for (i, e) in entries.iter().enumerate() {
        let text = batch_text(e)?;
        let lvl: EcLevel = e
            .level
            .as_deref()
            .map(|s| s.parse())
            .unwrap_or(Ok(EcLevel::M))
            .map_err(|e: QrError| anyhow::anyhow!(e))?;
        let cfg = QrConfig {
            level: lvl,
            version: VersionMode::Auto,
            margin: opts.margin,
        };
        let qr = QrCode::encode(&text, cfg).map_err(|e| anyhow::anyhow!(e))?;
        let path = format!("{out}/qr_{:04}.png", i + 1);
        std::fs::write(
            &path,
            qr.to_png_bytes(opts.size, None)
                .map_err(|e| anyhow::anyhow!("{e}"))?,
        )?;
        pb.set_message(path.clone());
        pb.inc(1);
    }
    pb.finish_with_message(format!("完成: {total} 个 QR → {out}"));
    Ok(())
}

fn batch_text(e: &BatchEntry) -> Result<String> {
    if let Some(c) = &e.content {
        return Ok(c.clone());
    }
    if let Some(u) = &e.url {
        return Ok(u.clone());
    }
    if let Some(s) = &e.ssid {
        return Ok(text_builder::build_wifi_text(
            s,
            e.password.as_deref().unwrap_or(""),
            e.encryption.as_deref().unwrap_or("WPA"),
            e.hidden.unwrap_or(false),
        ));
    }
    if let Some(n) = &e.name {
        return Ok(text_builder::build_vcard_text(
            n,
            e.phone.as_deref().unwrap_or(""),
            e.email.as_deref().unwrap_or(""),
            e.company.as_deref().unwrap_or(""),
            e.address.as_deref().unwrap_or(""),
            "",
            "",
            "",
            "",
            "",
        ));
    }
    if let Some(t) = &e.to {
        return Ok(text_builder::build_email_text(
            t,
            e.subject.as_deref().unwrap_or(""),
            e.body.as_deref().unwrap_or(""),
        ));
    }
    if let Some(n) = &e.number {
        if let Some(m) = &e.message {
            return Ok(text_builder::build_sms_text(n, m));
        }
        return Ok(text_builder::build_phone_text(n));
    }
    bail!("无法识别的条目格式")
}

// ──────────────────── 工具函数 ────────────────────

fn parse_csv(input: &str) -> Result<Vec<BatchEntry>, String> {
    let mut lines = input.lines();
    let cols: Vec<&str> = lines
        .next()
        .ok_or("CSV 为空")?
        .split(',')
        .map(|s| s.trim())
        .collect();
    let mut out = Vec::new();
    for line in lines {
        if line.trim().is_empty() {
            continue;
        }
        let values: Vec<String> = line
            .split(',')
            .map(|s| s.trim().trim_matches('"').to_string())
            .collect();

        // 使用 HashMap 替代 17 字段元组解构
        let mut map = std::collections::HashMap::new();
        for (i, col) in cols.iter().enumerate() {
            if let Some(val) = values.get(i).cloned() {
                map.insert(*col, val);
            }
        }

        let get = |k: &str| map.get(k).cloned();
        let hidden = get("hidden").map(|x| x == "true");

        out.push(BatchEntry {
            content: get("content"),
            level: get("level"),
            ssid: get("ssid"),
            password: get("password"),
            encryption: get("encryption"),
            hidden,
            name: get("name"),
            phone: get("phone"),
            email: get("email"),
            company: get("company"),
            address: get("address"),
            to: get("to"),
            subject: get("subject"),
            body: get("body"),
            number: get("number"),
            message: get("message"),
            url: get("url"),
        });
    }
    Ok(out)
}

fn check_path(p: &str) -> Result<()> {
    let path = Path::new(p);
    // 禁止绝对路径
    if path.is_absolute() {
        bail!("不允许使用绝对路径: {p}");
    }
    // 禁止父目录遍历
    if path
        .components()
        .any(|c| matches!(c, std::path::Component::ParentDir))
    {
        bail!("路径不允许包含 '..': {p}");
    }
    Ok(())
}