recover_pendrive.py

#!/usr/bin/env python3
"""
Safe wrapper around photorec for recovering files from a pendrive.
"""

from __future__ import annotations

import argparse
import os
import shutil
import stat
import subprocess
import sys
from pathlib import Path
from typing import Dict, Optional, Tuple


def ask_yes_no(question: str, default: bool = False) -> bool:
    prompt = "[Y/n]" if default else "[y/N]"
    while True:
        answer = input(f"{question} {prompt} ").strip().lower()
        if not answer:
            return default
        if answer in {"y", "yes"}:
            return True
        if answer in {"n", "no"}:
            return False
        print("Please answer yes or no.")


def read_mounts() -> Dict[str, str]:
    mounts: Dict[str, str] = {}
    proc_mounts = Path("/proc/mounts")
    if not proc_mounts.exists():
        return mounts

    with proc_mounts.open("r", encoding="utf-8") as fd:
        for line in fd:
            parts = line.split()
            if len(parts) >= 2:
                source = parts[0]
                target = parts[1].replace("\\040", " ")
                mounts[target] = source
    return mounts


def find_mount_source(path: Path, mounts: Dict[str, str]) -> Optional[Tuple[str, str]]:
    resolved = path.resolve()
    best_mount = None
    best_len = -1
    for mountpoint, source in mounts.items():
        mount_path = Path(mountpoint)
        try:
            resolved.relative_to(mount_path)
        except ValueError:
            continue
        mount_len = len(str(mount_path))
        if mount_len > best_len:
            best_len = mount_len
            best_mount = (mountpoint, source)
    return best_mount


def is_output_on_device(device: Path, output_dir: Path) -> bool:
    mounts = read_mounts()
    match = find_mount_source(output_dir, mounts)
    if not match:
        return False

    _, source = match
    device_str = str(device)

    if source == device_str:
        return True
    if source.startswith(f"{device_str}p"):  # nvme style
        return True
    if source.startswith(device_str) and source[len(device_str) :].isdigit():
        return True
    return False


def list_device_mountpoints(device: Path) -> Dict[str, str]:
    mounts = read_mounts()
    device_str = str(device)
    out = {}
    for mountpoint, source in mounts.items():
        if source == device_str:
            out[mountpoint] = source
        elif source.startswith(f"{device_str}p"):
            out[mountpoint] = source
        elif source.startswith(device_str) and source[len(device_str) :].isdigit():
            out[mountpoint] = source
    return out


def ensure_block_device(path: Path) -> None:
    if not path.exists():
        raise ValueError(f"Device path does not exist: {path}")

    mode = path.stat().st_mode
    if not stat.S_ISBLK(mode):
        raise ValueError(f"Path is not a block device: {path}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Recover deleted files from a pendrive (including formatted partitions) "
            "using foremost, a professional forensic file carving tool."
        )
    )
    parser.add_argument("--device", required=True, help="Block device path (e.g. /dev/sdb)")
    parser.add_argument(
        "--output",
        required=True,
        help="Local output folder where recovered files will be saved",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    foremost = shutil.which("foremost")
    if not foremost:
        print("Error: foremost not found. Install it with: sudo apt install foremost")
        return 1

    device = Path(args.device)
    output_dir = Path(args.output).expanduser().resolve()

    try:
        ensure_block_device(device)
    except ValueError as exc:
        print(f"Error: {exc}")
        return 1

    output_dir.mkdir(parents=True, exist_ok=True)

    if is_output_on_device(device, output_dir):
        print(
            "Error: output directory appears to be mounted from the same pendrive. "
            "Choose a local disk folder instead."
        )
        return 1

    mounted_parts = list_device_mountpoints(device)

    print("Recovery summary")
    print(f"- Source device: {device}")
    print(f"- Recovery destination: {output_dir}")
    print("- Engine: foremost (forensic file carving)")
    print("- Note: recovered files will be organized by type in subdirectories")
    print()

    if mounted_parts:
        print("Warning: the following partitions are currently mounted:")
        for mountpoint, source in sorted(mounted_parts.items()):
            print(f"  - {source} -> {mountpoint}")
        print("For best results, unmount pendrive partitions before recovery.")
        if not ask_yes_no("Do you want to continue anyway?", default=False):
            print("Cancelled.")
            return 0

    if os.geteuid() != 0:
        print("Warning: running as non-root may fail to read the device.")
        if not ask_yes_no("Continue without root permissions?", default=False):
            print("Cancelled.")
            return 0

    if not ask_yes_no(
        "Proceed with recovery now? (This will only write to the local output folder)",
        default=False,
    ):
        print("Cancelled.")
        return 0

    # Use foremost for recovery - more effective than photorec for formatted drives
    command = [foremost, "-v", "-i", str(device), "-o", str(output_dir)]
    print("\nLaunching foremost in unattended mode...")
    print(" ".join(command))
    print("Recovering all file types using forensic carving...\n")
    
    try:
        result = subprocess.run(
            command,
            input=automated_input.encode(),
            check=False
        )
    except KeyboardInterrupt:
        print("\nInterrupted by user.")
        return 130
        return 130

    if result.returncode == 0:
        print(f"\nRecovery process finished. Check recovered files under: {output_dir}")
        return 0

    print(
        f"\nphotorec exited with code {result.returncode}. "
        "Please review on-screen messages and retry."
    )
    return result.returncode


if __name__ == "__main__":
    sys.exit(main())