tasign 0.1.1

//! ELF 与 `.ta_signature` 节相关的**用户态**工具（`std`）。
//!
//! - **生产默认**：[`append_ta_signature`] / [`append_ta_signature_objcopy`] 使用 GNU `objcopy --add-section`，
//!   适用于 `ET_EXEC` / `ET_DYN` / `ET_REL` 等常见 ELF。
//! - **仅验证与实验**：[`append_ta_signature_object`]、[`append_ta_signature_bytes`] 基于 `object` crate 读写，
//!   **仅支持 relocatable（`ET_REL`）**；不作为 CLI `write-elf` 的默认实现。

use std::collections::HashMap;
use std::fs;
use std::path::Path;
use std::process::Command;

use object::{
    write, Object, ObjectComdat, ObjectKind, ObjectSection, ObjectSymbol, RelocationTarget,
    SectionFlags, SectionKind, SymbolFlags, SymbolKind, SymbolSection,
};

use crate::cms::sign_sm2_cms_with_artifacts;
use crate::error::Error;
use crate::plain::{build_plain_bin, TA_SIGNATURE_SECTION};
use crate::{CmsSignAlgorithm, SignInputs};

/// 写出中间文件：`plain.bin`、`signature.bin`，可选 `signed_attrs.der`（由调用方在 cms 层生成时另行保存）。
/// `gmssl_path` 为 `None` 时按 [`crate::gmssl_cli::resolve_gmssl_path`]（CLI 未传 `--gmssl` → `GMSSL` → PATH → 默认路径）解析。
pub fn write_outputs(
    elf_path: &Path,
    out_dir: &Path,
    leaf_cert_der: &[u8],
    intermediate_certs_der: &[Vec<u8>],
    leaf_key_path: &Path,
    leaf_key_pass: &str,
    gmssl_path: Option<&Path>,
    algorithm: CmsSignAlgorithm,
) -> Result<(), Error> {
    fs::create_dir_all(out_dir)?;
    let elf_bytes = fs::read(elf_path)?;
    let plain = build_plain_bin(&elf_bytes)?;
    fs::write(out_dir.join("plain.bin"), &plain)?;

    let gmssl_resolved = crate::gmssl_cli::resolve_gmssl_path(gmssl_path);
    let intermediate_refs: Vec<&[u8]> = intermediate_certs_der.iter().map(|v| v.as_slice()).collect();
    let out = sign_sm2_cms_with_artifacts(SignInputs {
        plain: &plain,
        leaf_cert_der,
        intermediate_certs_der: &intermediate_refs,
        cms_attached: false,
        cms_use_gmssl_oid: false,
        leaf_key_path,
        leaf_key_pass,
        gmssl_path: Some(gmssl_resolved.as_path()),
        algorithm,
    })?;
    fs::write(out_dir.join("signed_attrs.der"), out.signed_attrs_der)?;
    fs::write(out_dir.join("signature.bin"), out.signature_der)?;
    Ok(())
}

/// 将 PKCS#7 DER 作为 `.ta_signature` 节追加到 ELF，写出 `output_elf`（默认使用 `objcopy`）。
pub fn append_ta_signature(
    input_elf: &Path,
    signature_bin: &Path,
    output_elf: &Path,
) -> Result<(), Error> {
    append_ta_signature_objcopy(input_elf, signature_bin, output_elf, "objcopy")
}

/// 使用 `objcopy` 将 PKCS#7 DER 作为 `.ta_signature` 节追加到 ELF。
///
/// 要求 `.ta_signature` 节标志为 `noload`。
pub fn append_ta_signature_objcopy(
    input_elf: &Path,
    signature_bin: &Path,
    output_elf: &Path,
    objcopy: &str,
) -> Result<(), Error> {
    let section_arg = format!("{TA_SIGNATURE_SECTION}={}", signature_bin.to_string_lossy());
    let flags_arg = format!("{TA_SIGNATURE_SECTION}=noload");
    let st = Command::new(objcopy)
        .args([
            "--add-section",
            &section_arg,
            "--set-section-flags",
            &flags_arg,
            &input_elf.to_string_lossy(),
            &output_elf.to_string_lossy(),
        ])
        .status()
        .map_err(Error::Io)?;
    if !st.success() {
        return Err(Error::ElfWrite(format!(
            "{objcopy} --add-section/--set-section-flags failed: {st}"
        )));
    }
    Ok(())
}

/// 将 PKCS#7 DER 作为 `.ta_signature` 节追加到 relocatable ELF（`ET_REL`），写出 `output_elf`。
///
/// 该函数保留用于 object 写路径验证，不是 CLI 默认路径。
pub fn append_ta_signature_object(
    input_elf: &Path,
    signature_bin: &Path,
    output_elf: &Path,
) -> Result<(), Error> {
    let elf_bytes = fs::read(input_elf)?;
    let sig = fs::read(signature_bin)?;
    let out = append_ta_signature_bytes(&elf_bytes, &sig)?;
    fs::write(output_elf, &out)?;
    Ok(())
}

/// 将 relocatable ELF 经 `object` read→write 归一化（布局与 [`append_ta_signature_bytes`] 一致）。
///
/// 当前 `plain.bin` 协议以 goblin 直接解析**原始** ELF 字节为准（见 `docs/plain_protocol.md`）。
/// 本函数仅用于需要 `object` 规范布局的测试或迁移场景，**不要**假定签名流程必须先归一化再 `build_plain_bin`。
pub fn normalize_relocatable_elf_bytes(elf_bytes: &[u8]) -> Result<Vec<u8>, Error> {
    relocatable_copy_elf(
        elf_bytes,
        RelocElfOpts {
            skip_ta_signature: false,
            append_ta_signature: None,
        },
    )
}

/// 去掉 `.ta_signature` 节并写出 relocatable ELF（与未嵌入签名时的归一化布局一致）。
pub fn strip_ta_signature_relocatable_bytes(elf_bytes: &[u8]) -> Result<Vec<u8>, Error> {
    let normalized = normalize_relocatable_elf_bytes(elf_bytes)?;
    relocatable_copy_elf(
        &normalized,
        RelocElfOpts {
            skip_ta_signature: true,
            append_ta_signature: None,
        },
    )
}

/// 将 PKCS#7 DER 作为 `.ta_signature` 节追加到 relocatable ELF 字节，返回新 ELF。
///
/// 直接基于输入 ELF 追加 `.ta_signature`，不改变现有 section 过滤协议语义。
pub fn append_ta_signature_bytes(elf_bytes: &[u8], signature_der: &[u8]) -> Result<Vec<u8>, Error> {
    relocatable_copy_elf(
        elf_bytes,
        RelocElfOpts {
            skip_ta_signature: false,
            append_ta_signature: Some(signature_der),
        },
    )
}

struct RelocElfOpts<'a> {
    skip_ta_signature: bool,
    append_ta_signature: Option<&'a [u8]>,
}

fn relocatable_copy_elf(in_data: &[u8], opts: RelocElfOpts<'_>) -> Result<Vec<u8>, Error> {
    let in_object = object::File::parse(in_data).map_err(|e| Error::ElfWrite(e.to_string()))?;
    if in_object.kind() != ObjectKind::Relocatable {
        return Err(Error::ElfWrite(format!(
            "仅支持 relocatable 对象文件（ET_REL），当前: {:?}",
            in_object.kind()
        )));
    }

    let mut out_object = write::Object::new(
        in_object.format(),
        in_object.architecture(),
        in_object.endianness(),
    );
    out_object.mangling = write::Mangling::None;
    out_object.flags = in_object.flags();

    let mut out_sections: HashMap<object::SectionIndex, write::SectionId> = HashMap::new();

    for in_section in in_object.sections() {
        if in_section.kind() == SectionKind::Metadata {
            continue;
        }
        let name = in_section.name().unwrap_or("");
        if opts.skip_ta_signature && name == TA_SIGNATURE_SECTION {
            continue;
        }

        let section_id = out_object.add_section(
            in_section
                .segment_name()
                .map_err(|e| Error::ElfWrite(e.to_string()))?
                .unwrap_or("")
                .as_bytes()
                .to_vec(),
            in_section
                .name()
                .map_err(|e| Error::ElfWrite(e.to_string()))?
                .as_bytes()
                .to_vec(),
            in_section.kind(),
        );
        let out_section = out_object.section_mut(section_id);
        if out_section.is_bss() {
            out_section.append_bss(in_section.size(), in_section.align());
        } else {
            let data = in_section
                .data()
                .map_err(|e| Error::ElfWrite(e.to_string()))?;
            out_section.set_data(data, in_section.align());
        }
        out_section.flags = in_section.flags();
        out_sections.insert(in_section.index(), section_id);
    }

    if let Some(data) = opts.append_ta_signature {
        let section_id = out_object.add_section(
            vec![],
            TA_SIGNATURE_SECTION.as_bytes().to_vec(),
            SectionKind::Data,
        );
        let out_section = out_object.section_mut(section_id);
        out_section.set_data(data, 1);
        out_section.flags = SectionFlags::Elf {
            sh_flags: (object::elf::SHF_ALLOC | object::elf::SHF_WRITE) as u64,
        };
    }

    let mut out_symbols: HashMap<object::SymbolIndex, write::SymbolId> = HashMap::new();
    for in_symbol in in_object.symbols() {
        let (section, value) = match in_symbol.section() {
            SymbolSection::None => (write::SymbolSection::None, in_symbol.address()),
            SymbolSection::Undefined => (write::SymbolSection::Undefined, in_symbol.address()),
            SymbolSection::Absolute => (write::SymbolSection::Absolute, in_symbol.address()),
            SymbolSection::Common => (write::SymbolSection::Common, in_symbol.address()),
            SymbolSection::Section(index) => {
                if let Some(out_section) = out_sections.get(&index) {
                    (
                        write::SymbolSection::Section(*out_section),
                        in_symbol.address()
                            - in_object
                                .section_by_index(index)
                                .map_err(|e| Error::ElfWrite(e.to_string()))?
                                .address(),
                    )
                } else {
                    if in_symbol.kind() != SymbolKind::Section {
                        return Err(Error::ElfWrite(format!(
                            "symbol references skipped section: {:?}",
                            in_symbol
                        )));
                    }
                    continue;
                }
            }
            _ => {
                return Err(Error::ElfWrite(format!(
                    "unsupported symbol section for {:?}",
                    in_symbol
                )));
            }
        };
        let flags = match in_symbol.flags() {
            SymbolFlags::None => SymbolFlags::None,
            SymbolFlags::Elf { st_info, st_other } => SymbolFlags::Elf { st_info, st_other },
            SymbolFlags::MachO { n_desc } => SymbolFlags::MachO { n_desc },
            SymbolFlags::CoffSection {
                selection,
                associative_section,
            } => {
                let associative_section = match associative_section {
                    None => None,
                    Some(index) => Some(
                        *out_sections.get(&index).ok_or_else(|| {
                            Error::ElfWrite("CoffSection associative_section index missing".into())
                        })?,
                    ),
                };
                SymbolFlags::CoffSection {
                    selection,
                    associative_section,
                }
            }
            SymbolFlags::Xcoff {
                n_sclass,
                x_smtyp,
                x_smclas,
                containing_csect,
            } => {
                let containing_csect = match containing_csect {
                    None => None,
                    Some(index) => Some(
                        *out_symbols.get(&index).ok_or_else(|| {
                            Error::ElfWrite("Xcoff containing_csect symbol missing".into())
                        })?,
                    ),
                };
                SymbolFlags::Xcoff {
                    n_sclass,
                    x_smtyp,
                    x_smclas,
                    containing_csect,
                }
            }
            _ => {
                return Err(Error::ElfWrite(format!(
                    "unsupported symbol flags for {:?}",
                    in_symbol
                )));
            }
        };
        let out_symbol = write::Symbol {
            name: in_symbol.name().unwrap_or("").as_bytes().to_vec(),
            value,
            size: in_symbol.size(),
            kind: in_symbol.kind(),
            scope: in_symbol.scope(),
            weak: in_symbol.is_weak(),
            section,
            flags,
        };
        let symbol_id = out_object.add_symbol(out_symbol);
        out_symbols.insert(in_symbol.index(), symbol_id);
    }

    for in_section in in_object.sections() {
        if in_section.kind() == SectionKind::Metadata {
            continue;
        }
        let name = in_section.name().unwrap_or("");
        if opts.skip_ta_signature && name == TA_SIGNATURE_SECTION {
            continue;
        }
        let out_section = *out_sections
            .get(&in_section.index())
            .ok_or_else(|| Error::ElfWrite("missing out section for relocation pass".into()))?;
        for (offset, in_relocation) in in_section.relocations() {
            let symbol = match in_relocation.target() {
                RelocationTarget::Symbol(symbol) => *out_symbols.get(&symbol).ok_or_else(|| {
                    Error::ElfWrite("relocation symbol index missing in out_symbols".into())
                })?,
                RelocationTarget::Section(section) => out_object.section_symbol(
                    *out_sections.get(&section).ok_or_else(|| {
                        Error::ElfWrite("relocation section index missing".into())
                    })?,
                ),
                _ => {
                    return Err(Error::ElfWrite(format!(
                        "unsupported relocation target: {:?}",
                        in_relocation
                    )));
                }
            };
            let out_relocation = write::Relocation {
                offset,
                symbol,
                addend: in_relocation.addend(),
                flags: in_relocation.flags(),
            };
            out_object
                .add_relocation(out_section, out_relocation)
                .map_err(|e| Error::ElfWrite(e.to_string()))?;
        }
    }

    for in_comdat in in_object.comdats() {
        let mut sections = Vec::new();
        for in_section in in_comdat.sections() {
            sections.push(
                *out_sections
                    .get(&in_section)
                    .ok_or_else(|| Error::ElfWrite("comdat section missing".into()))?,
            );
        }
        out_object.add_comdat(write::Comdat {
            kind: in_comdat.kind(),
            symbol: *out_symbols
                .get(&in_comdat.symbol())
                .ok_or_else(|| Error::ElfWrite("comdat symbol missing".into()))?,
            sections,
        });
    }

    if let Some(in_build_version) = match &in_object {
        object::File::MachO32(file) => file.build_version().unwrap(),
        object::File::MachO64(file) => file.build_version().unwrap(),
        _ => None,
    } {
        let mut out_build_version = object::write::MachOBuildVersion::default();
        out_build_version.platform = in_build_version.platform.get(in_object.endianness());
        out_build_version.minos = in_build_version.minos.get(in_object.endianness());
        out_build_version.sdk = in_build_version.sdk.get(in_object.endianness());
        out_object.set_macho_build_version(out_build_version);
    }

    out_object
        .write()
        .map_err(|e| Error::ElfWrite(e.to_string()))
}