#!/usr/bin/env python3
"""p7s/p7m‑extract – unwrap CMS (PKCS#7) SignedData files to their payload.

Handles typical e‑mail attachments:
  • **.p7s** – detached signature container (often spam.pdf.p7s)
  • **.p7m** – opaque/mime SignedData wrapper (often complete document)

The script calls **OpenSSL** once per file and writes the decoded payload next
to – or inside a chosen folder – preserving original filenames whenever
possible.

Features
--------
* Auto‑detects the *openssl* executable (or use --openssl=PATH).
* Multiprocessing via *concurrent.futures* (defaults to CPU count).
* Safe filename derivation ("*.pdf.p7s" → "*.pdf"; "foo.p7m" → "foo.bin").
* Verbose logging with -v / -vv, progress bar when *tqdm* present.
* Non‑zero exit code if any extraction fails.
"""
from __future__ import annotations

import argparse
import logging
import os
import shutil
import subprocess
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence

try:
    from tqdm import tqdm  # type: ignore
except ImportError:  # pragma: no cover
    tqdm = None  # type: ignore

LOGGER = logging.getLogger("p7extract")
SUPPORTED_SUFFIXES: tuple[str, ...] = (".p7s", ".p7m")


@dataclass(slots=True)
class Task:
    src: Path
    dst: Path


# ---------------------------------------------------------------------------
# Utility functions
# ---------------------------------------------------------------------------

def find_openssl(explicit: str | None = None) -> str:
    """Return an *openssl* executable path or raise FileNotFoundError."""
    if explicit:
        if shutil.which(explicit):
            return explicit
        raise FileNotFoundError(f"Specified openssl not found: {explicit}")
    exe = shutil.which("openssl")
    if exe:
        return exe
    raise FileNotFoundError("'openssl' executable not found – install it or pass --openssl")


def derive_dst(src: Path) -> Path:
    """Return destination Path derived from *src* (.p7s / .p7m stripped)."""
    if src.suffix.lower() not in SUPPORTED_SUFFIXES:
        return src.with_suffix(".bin")
    stem = src.stem  # drop .p7s/.p7m
    if Path(stem).suffix:
        return src.with_suffix("")  # keep existing extension (e.g. report.pdf)
    return src.with_suffix(".bin")  # otherwise default to .bin


def collect_tasks(root: Path, out_dir: Path, recursive: bool) -> list[Task]:
    """Return Task list for every .p7s / .p7m under *root*."""
    tasks: list[Task] = []
    iterator = root.rglob if recursive else root.glob
    for suf in SUPPORTED_SUFFIXES:
        for src in iterator(f"*{suf}"):
            if not src.is_file():
                continue
            dst = out_dir / derive_dst(src)
            dst.parent.mkdir(parents=True, exist_ok=True)
            tasks.append(Task(src, dst))
    return tasks


# ---------------------------------------------------------------------------
# Core extraction logic (multiprocess‑safe)
# ---------------------------------------------------------------------------

def openssl_extract(task: Task, openssl_cmd: str) -> tuple[Task, bool, str | None]:
    """Run *openssl cms -verify* for *task*.  Return (task, ok, errmsg)."""
    cmd = [
        openssl_cmd,
        "cms",
        "-verify",
        "-noverify",
        "-nocerts",
        "-nosigs",
        "-binary",
        "-inform",
        "DER",
        "-in",
        str(task.src),
        "-out",
        str(task.dst),
    ]
    try:
        res = subprocess.run(cmd, capture_output=True, text=True)
    except Exception as exc:  # pragma: no cover – unexpected runtime failure
        return task, False, str(exc)

    if res.returncode:
        return task, False, res.stderr.strip() or "openssl returned non‑zero"

    if not task.dst.exists() or task.dst.stat().st_size == 0:
        try:
            task.dst.unlink(missing_ok=True)
        except FileNotFoundError:
            pass
        return task, False, "no payload extracted (detached signature?)"

    return task, True, None


# ---------------------------------------------------------------------------
# CLI / main
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Extract payloads from .p7s / .p7m files using OpenSSL")
    p.add_argument("input", type=Path, help="Directory containing .p7s / .p7m files")
    p.add_argument("output", type=Path, nargs="?", help="Destination directory (default: same as input)")
    p.add_argument("-r", "--recursive", action="store_true", help="Search sub‑directories recursively")
    p.add_argument("-w", "--workers", type=int, help="Number of parallel processes (default: CPU count)")
    p.add_argument("--openssl", help="Path to openssl executable")
    p.add_argument("-v", "--verbose", action="count", default=0, help="Increase verbosity (‑v / ‑vv)")
    return p


def configure_logging(verbosity: int) -> None:
    level = logging.WARNING - min(verbosity, 2) * 10
    logging.basicConfig(level=level, format="%(levelname)s: %(message)s")


def main(argv: Sequence[str] | None = None) -> None:  # noqa: D401 – imperative
    args = build_parser().parse_args(argv)
    configure_logging(args.verbose)

    try:
        openssl_cmd = find_openssl(args.openssl)
    except FileNotFoundError as exc:
        LOGGER.error(exc)
        sys.exit(10)

    input_dir = args.input.expanduser().resolve()
    if not input_dir.is_dir():
        LOGGER.error("%s is not a directory", input_dir)
        sys.exit(10)

    output_dir = (args.output or input_dir).expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    tasks = collect_tasks(input_dir, output_dir, args.recursive)
    if not tasks:
        LOGGER.warning("No .p7s / .p7m files found in %s", input_dir)
        sys.exit(2)

    LOGGER.info("Found %d file(s) to extract", len(tasks))

    workers = args.workers or os.cpu_count() or 1
    LOGGER.debug("Using %d workers", workers)

    bar = None
    if tqdm is not None and sys.stdout.isatty():
        bar = tqdm(total=len(tasks), desc="Extracting", unit="file")  # type: ignore

    ok = fail = 0
    with ProcessPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(openssl_extract, t, openssl_cmd): t for t in tasks}
        for fut in as_completed(futures):
            task, success, err = fut.result()
            if success:
                LOGGER.info("✔ %s → %s", task.src.name, task.dst.name)
                ok += 1
            else:
                LOGGER.error("✖ %s: %s", task.src.name, err)
                fail += 1
            if bar:
                bar.update(1)
    if bar:
        bar.close()

    LOGGER.info("Done – %d succeeded, %d failed", ok, fail)
    sys.exit(0 if fail == 0 else 1)


if __name__ == "__main__":  # pragma: no cover
    main()
Any mistakes?