fcmc 0.6.3

use bollard::Docker;
use clap::Parser;
use colored::*;
use fcmc::ChallengeMeta;
use std::{
    fs,
    io::{self, Write},
    path::PathBuf,
};
#[derive(Parser, Debug)]
#[command(name = "fcmc", about = "FloatCTF 题目配置检查和管理工具")]
#[command(version = "0.6.0")]
struct Args {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Parser, Debug, Clone)]
#[command(rename_all = "snake_case")]
enum Commands {
    /// 检查题目配置文件是否合法
    Check {
        /// 配置文件目录 (里面需要包含 meta.toml)
        #[arg(short, long)]
        path: Option<String>,
    },
    /// 构建题目镜像
    Build {
        /// 配置文件目录 (里面需要包含 meta.toml)
        #[arg(short, long)]
        path: Option<String>,
    },
    /// 生成新的题目模板
    Gen {
        /// 新题目的名称
        #[arg(short, long)]
        name: String,

        /// 输出目录
        #[arg(short, long, default_value = ".")]
        output: String,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    match args.command {
        Commands::Check { path } => {
            // 如果指定了 --file，用它作为目录；否则默认当前目录
            let mut pass = true;

            let dir = path.unwrap_or_else(|| ".".to_string());
            let path = PathBuf::from(&dir).join("meta.toml");

            println!("\n================ 配置检查报告 ================\n");
            println!("配置文件: {:?}", path);

            let content = fs::read_to_string(&path)?;
            let cfg = toml::from_str::<ChallengeMeta>(&content);
            //  检查配置文件是否合法
            println!("\n[解析结果]");
            match cfg {
                Ok(cfg) => {
                    println!("  {}   配置文件解析成功", "OK".green());

                    println!("\n[附件检查]");
                    // 检查附件
                    if let Some(attachment) = &cfg.attachment {
                        let attachment_path = PathBuf::from(&dir).join(attachment);
                        if attachment_path.exists() {
                            let size = fs::metadata(&attachment_path)?.len();
                            println!(
                                "  {}   附件存在: {:?} ({} bytes)",
                                "OK".green(),
                                attachment_path,
                                size
                            );
                        } else {
                            println!("  {}   附件不存在: {:?}", "ERR".red(), attachment_path);
                            pass = false;
                        }
                    } else {
                        println!("  {}   未配置附件", "WARN".yellow());
                    }

                    // test docker
                    println!("\n[Docker 检查]");
                    if cfg.docker.is_some() {
                        match test_docker(&cfg).await {
                            Ok(_) => (),
                            Err(e) => {
                                println!("  {}   Docker 测试失败: {}", "ERR".red(), e);
                                println!("{} fcmc build -p {}", "  尝试构建镜像:".cyan(), dir);
                                pass = false;
                            }
                        }
                    } else {
                        println!("  {}   未配置 Docker", "WARN".yellow());
                    }
                }
                Err(e) => {
                    println!("  {}   配置文件解析失败: {}", "ERR".red(), e);
                    pass = false;
                }
            }

            if pass {
                println!("\n----------------------------------------------");
                println!("最终结果: {}", "通过".green());
                println!("==============================================\n");
            } else {
                println!("\n----------------------------------------------");
                println!("最终结果: {}", "失败".red());
                println!("==============================================\n");
            }
        }
        Commands::Gen { name, output } => {
            generate_template(&name, &output).await?;
        }
        Commands::Build { path } => {
            let dir = path.unwrap_or_else(|| ".".to_string());
            let path = PathBuf::from(&dir).join("meta.toml");

            let content = fs::read_to_string(&path)?;
            let cfg = toml::from_str::<ChallengeMeta>(&content)?;

            let src_dir = PathBuf::from(&dir).join("src");
            build_image(&cfg, &src_dir).await?;
        }
    }

    Ok(())
}

async fn test_docker(cm: &ChallengeMeta) -> anyhow::Result<()> {
    let docker = Docker::connect_with_defaults()?;

    let flag = "flag{test-test-test-test}";

    match cm.create_and_start(&docker, "local_test", flag).await {
        Ok(port) => {
            let name = format!("{}_local_test", cm.name);
            println!("  {} 已启动容器: {}", "OK".green(), name);
            println!("  {} 访问地址: http://localhost:{}", "INFO".cyan(), port);

            // kill
            print!("  INFO 按回车键继续以清理容器...");
            io::stdout().flush()?;
            let mut input = String::new();
            io::stdin().read_line(&mut input)?;

            match fcmc::stop_and_remove(&docker, "local_test").await {
                Ok(_) => println!("  {} 已清理容器: {}", "OK".green(), name),
                Err(e) => println!("  {} 容器清理失败: {}", "ERR".red(), e),
            }
        }
        Err(e) => {
            return Err(e);
        }
    }

    Ok(())
}

async fn generate_template(name: &str, output_dir: &str) -> anyhow::Result<()> {
    use std::fs;
    use std::path::Path;

    // 创建挑战目录
    let challenge_dir = Path::new(output_dir).join(name);
    fs::create_dir_all(&challenge_dir)?;

    // 创建src目录
    let src_dir = challenge_dir.join("src");
    fs::create_dir_all(&src_dir)?;

    // 创建attachment目录
    let attachment_dir = challenge_dir.join("attachment");
    fs::create_dir_all(&attachment_dir)?;

    // 创建meta.toml文件
    let meta_content = format!(
        r#"name = "{}"
author = "your_email@example.com" # modify
category = "Web" # modify
description = "Challenge description" # modify

attachment = "attachment/src.zip" # Optional

[flag]
value = ""       # is empty stand for dynamic flag # modify
env_var = "FLAG"


[docker]
image_tag = "floatctf/name:challenge-web_v1.0" # modify
port = "80/tcp"

"#,
        name
    );
    fs::write(challenge_dir.join("meta.toml"), meta_content)?;

    // 创建flag
    let flag_content = "flag{test_flag}";
    fs::write(src_dir.join("flag"), flag_content)?;

    // 创建flag.sh
    let flag_sh_content = r#"#!/bin/bash
# flag 动态替换脚本
sed -i "s/flag{test_flag}/$FLAG/" /flag

export FLAG=not_flag
FLAG=not_flag

rm -f /flag.sh
"#;
    fs::write(src_dir.join("flag.sh"), flag_sh_content)?;

    // 创建entrypoint.sh
    let entrypoint_content = r#"#!/bin/bash
if [ -f /flag.sh ]; then
    echo "--- 正在初始化 Flag ---"
    sed -i 's/\r//g' /flag.sh
    /flag.sh
fi

exec "$@"
"#;
    fs::write(src_dir.join("entrypoint.sh"), entrypoint_content)?;

    let index_php_content = r#"<?php
echo get_file_contents("/flag");
?>
"#;
    fs::write(src_dir.join("index.php"), index_php_content)?;
    // 创建Dockerfile
    let dockerfile_content = r#"FROM php:5-apache-jessie
LABEL Author="your_name <your_email@example.com>"

COPY flag /flag
COPY flag.sh /flag.sh
COPY entrypoint.sh /entrypoint.sh
COPY index.php /var/www/html/index.php
RUN chmod +x /flag.sh
RUN chmod +x /entrypoint.sh

# 必须
EXPOSE 80
WORKDIR /var/www/html


ENTRYPOINT [ "/entrypoint.sh" ]
CMD [ "apache2-foreground" ]
"#;
    fs::write(src_dir.join("Dockerfile"), dockerfile_content)?;

    println!("  {}   成功生成挑战模板: {:?}", "OK".green(), challenge_dir);
    println!(
        "  {}   请记得修改meta.toml中的作者和描述信息",
        "INFO".cyan()
    );
    println!(
        "  {}   请记得修改Dockerfile和源代码以实现你的挑战逻辑",
        "INFO".cyan()
    );

    Ok(())
}

async fn build_image(cm: &ChallengeMeta, src_dir: &PathBuf) -> anyhow::Result<()> {
    let docker = Docker::connect_with_defaults()?;
    let infos = cm.build_image(&docker, src_dir).await?;
    for info in infos {
        println!("{}", info);
    }
    Ok(())
}