beerpsi_fsdecrypt/src/main.rs

use std::{
    collections::HashMap,
    fs::{create_dir_all, File, FileTimes},
    io::{BufRead, BufReader, BufWriter, Write},
    path::{Path, PathBuf},
    time::{Duration, SystemTime},
};

use chrono::{FixedOffset, TimeZone};
use clap::Parser;
use exfat_fs::dir::{entry::fs::FsElement, Root};

use anyhow::{anyhow, Result};
use bootid::ContainerType;
use indicatif::{ProgressBar, ProgressStyle};
use ntfs::{
    indexes::NtfsFileNameIndex, structured_values::NtfsStandardInformation, Ntfs,
    NtfsAttributeType, NtfsTime,
};

use crate::stream::FscryptDecryptor;

mod bootid;
mod crypto;
mod stream;

/// Info collected from each input file for sorting and post-extraction merge.
struct InputFile {
    path: PathBuf,
    sequence_number: u8,
}

/// Info about an extracted VHD, used for delta merge.
struct ExtractedVhd {
    vhd_path: PathBuf,
    /// Original input .app file path (for resolving sibling files)
    input_path: PathBuf,
    sequence_number: u8,
    game_id: String,
}

fn exfat_timestamp_to_system_time(
    timestamp: &exfat_fs::timestamp::Timestamp,
) -> Result<SystemTime> {
    let exfat_date = timestamp.date();
    let exfat_time = timestamp.time();
    // exFAT UTC offset is in 15-minute intervals, so 1 = UTC+00:15, 2 = UTC+00:30, etc.
    let exfat_utc_offset = timestamp.utc_offset() as i32 * 15 * 60;
    let chrono_date_time = FixedOffset::east_opt(exfat_utc_offset)
        .ok_or_else(|| anyhow!("invaid utc offset: {}", timestamp.utc_offset()))?
        .with_ymd_and_hms(
            exfat_date.year as i32,
            exfat_date.month as u32,
            exfat_date.day as u32,
            exfat_time.hour as u32,
            exfat_time.minute as u32,
            exfat_time.second as u32,
        )
        .unwrap();

    return Ok(SystemTime::UNIX_EPOCH
        + Duration::from_micros(chrono_date_time.timestamp_micros().try_into()?));
}

fn extract_exfat_contents(exfat_path: &Path) -> Result<()> {
    // Create output directory with same name as exfat file (without extension)
    let output_dir = exfat_path.with_extension("");

    let file = FscryptDecryptor::new(File::open(exfat_path)?).map_err(|e| anyhow!(e))?;
    let mut root = Root::open(file)?;

    let pb = ProgressBar::new(calculate_exfat_size(root.items())?)
        .with_style(
            ProgressStyle::default_bar().template(
                "{prefix} [{bar:20!.bright.yellow/dim.white}] {bytes:>8} [{elapsed}<{eta}, {bytes_per_sec}]",
            )?
        );
    pb.set_prefix(format!(
        "Extracting {}",
        exfat_path.file_name().unwrap().display(),
    ));

    create_dir_all(&output_dir)?;
    extract_exfat_elements(root.items(), &output_dir, &pb)?;

    pb.finish();

    Ok(())
}

fn calculate_exfat_size(elements: &[FsElement<FscryptDecryptor<File>>]) -> Result<u64> {
    let mut total = 0;

    for element in elements {
        match element {
            FsElement::F(ref file) => total += file.len(),
            FsElement::D(directory) => total += calculate_exfat_size(&directory.open()?)?,
        }
    }

    Ok(total)
}

fn extract_exfat_elements(
    elements: &mut [FsElement<FscryptDecryptor<File>>],
    output_dir: &Path,
    pb: &ProgressBar,
) -> Result<()> {
    for element in elements {
        match element {
            FsElement::F(ref mut file) => {
                let dest_path = output_dir.join(file.name());
                let mut dest = File::create(dest_path)?;

                dest.set_times(
                    FileTimes::new()
                        .set_accessed(exfat_timestamp_to_system_time(
                            file.timestamps().accessed(),
                        )?)
                        .set_modified(exfat_timestamp_to_system_time(
                            file.timestamps().modified(),
                        )?),
                )?;

                let mut writer = BufWriter::with_capacity(256 * 1024, &mut dest);

                std::io::copy(file, &mut writer)?;
                writer.flush()?;
                pb.inc(file.len());
            }
            FsElement::D(directory) => {
                let dest_path = output_dir.join(directory.name());
                create_dir_all(&dest_path)?;

                let mut children = directory.open()?;
                extract_exfat_elements(&mut children, &dest_path, &pb)?;
            }
        }
    }

    Ok(())
}

fn ntfs_time_to_system_time(ntfs_time: NtfsTime) -> SystemTime {
    // An NTFS "interval" is 100 nanoseconds.
    // The Windows epoch is 1601-01-01, while the Unix epoch is 1970-01-01.
    let intervals_since_windows_epoch = ntfs_time.nt_timestamp();
    let intervals_since_unix_epoch = intervals_since_windows_epoch - 116_444_736_000_000_000;
    let nanos_since_unix_epoch = intervals_since_unix_epoch * 100;

    return SystemTime::UNIX_EPOCH + Duration::from_nanos(nanos_since_unix_epoch);
}

fn extract_internal_vhd(image_path: &Path, sequence_number: u8) -> Result<PathBuf> {
    let vhd_filename = format!("internal_{sequence_number}.vhd");
    let output_path = image_path.with_extension("vhd");

    let mut fs = FscryptDecryptor::new(File::open(image_path)?).map_err(|e| anyhow!(e))?;

    let mut ntfs = Ntfs::new(&mut fs)?;

    ntfs.read_upcase_table(&mut fs)?;

    let root_directory = ntfs.root_directory(&mut fs)?;
    let index = root_directory.directory_index(&mut fs)?;
    let mut finder = index.finder();
    let entry = NtfsFileNameIndex::find(&mut finder, &ntfs, &mut fs, &vhd_filename)
        .ok_or_else(|| anyhow!("could not find VHD {vhd_filename}"))??;
    let file = entry.to_file(&ntfs, &mut fs)?;
    let data_item = file
        .data(&mut fs, "")
        .ok_or_else(|| anyhow!("file data does not exist"))??;
    let data_attribute = data_item.to_attribute()?;
    let mut data_value =
        BufReader::with_capacity(256 * 1024, data_attribute.value(&mut fs)?.attach(&mut fs));

    let mut output_file = File::create(&output_path)?;

    let pb = ProgressBar::new(data_attribute.value_length() as u64)
        .with_style(
            ProgressStyle::default_bar()
                .template("{prefix} [{bar:20!.bright.yellow/dim.white}] {bytes:>8} [{elapsed}<{eta}, {bytes_per_sec}]")?
        );
    pb.set_prefix(format!("{}", output_path.file_name().unwrap().display()));

    loop {
        let buffer = data_value.fill_buf()?;
        let length = buffer.len();

        if length == 0 {
            break;
        }

        output_file.write_all(buffer)?;
        data_value.consume(length);
        pb.inc(length as u64);
    }
    output_file.flush()?;

    pb.finish();

    let mut attributes_iterator = file.attributes();

    while let Some(attribute) = attributes_iterator.next(&mut fs) {
        let attribute = attribute?;
        let attribute = attribute.to_attribute()?;

        match attribute.ty() {
            Ok(NtfsAttributeType::StandardInformation) => {
                let info = attribute.resident_structured_value::<NtfsStandardInformation>()?;

                output_file.set_times(
                    FileTimes::new()
                        .set_accessed(ntfs_time_to_system_time(info.access_time()))
                        .set_modified(ntfs_time_to_system_time(info.modification_time())),
                )?;

                break;
            }
            _ => continue,
        }
    }

    Ok(output_path)
}

/// Strip the `\\?\` extended-path prefix that `canonicalize` adds on Windows.
fn strip_extended_path_prefix(path: PathBuf) -> PathBuf {
    let s = path.to_string_lossy();
    if let Some(stripped) = s.strip_prefix(r"\\?\") {
        PathBuf::from(stripped)
    } else {
        path
    }
}

/// Mount a VHD, extract all contents to a folder with the same name, dismount, and delete the VHD.
/// Requires elevation (triggers UAC prompt).
#[cfg(windows)]
fn extract_vhd_contents(vhd_path: &Path) -> Result<()> {
    let vhd_abs = strip_extended_path_prefix(std::fs::canonicalize(vhd_path)?);
    let output_dir = vhd_path.with_extension("");
    let output_dir_abs = strip_extended_path_prefix(std::path::absolute(&output_dir)?);

    println!("Extracting VHD contents...");
    println!("  VHD:    {}", vhd_abs.display());
    println!("  Output: {}", output_dir_abs.display());

    // Write a PowerShell script to a temp file to avoid nested quoting hell.
    // Mount VHD read-only, copy contents, dismount.
    let script_path = vhd_path.with_extension("extract.ps1");
    let script_abs = strip_extended_path_prefix(std::path::absolute(&script_path)?);
    let error_log = vhd_path.with_extension("extract_error.txt");
    let error_log_abs = strip_extended_path_prefix(std::path::absolute(&error_log)?);

    let script = format!(
        r#"try {{
    $vhd = Mount-VHD -Path '{vhd}' -ReadOnly -Passthru
    $disk = $vhd | Get-Disk
    $part = $disk | Get-Partition | Where-Object {{ $_.Type -ne 'Reserved' }} | Select-Object -First 1

    # Try to get existing drive letter
    $dl = ($part | Get-Volume).DriveLetter

    # If no drive letter, assign one temporarily
    $assignedLetter = $false
    if (-not $dl) {{
        $usedLetters = (Get-Volume).DriveLetter
        $dl = [char[]](90..68) | Where-Object {{ $_ -notin $usedLetters }} | Select-Object -First 1
        if (-not $dl) {{
            Dismount-VHD -Path '{vhd}'
            throw 'No available drive letter'
        }}
        $part | Set-Partition -NewDriveLetter $dl
        $assignedLetter = $true
    }}

    $src = "$($dl):\"
    New-Item -ItemType Directory -Force -Path '{out}' | Out-Null
    Get-ChildItem -Path $src -Force | Where-Object {{ $_.Name -ne 'System Volume Information' -and $_.Name -ne '$Recycle.Bin' }} | Copy-Item -Destination '{out}' -Recurse -Force

    # Remove assigned letter before dismounting
    if ($assignedLetter) {{
        $part | Remove-PartitionAccessPath -AccessPath "$($dl):\" -ErrorAction SilentlyContinue
    }}
    Dismount-VHD -Path '{vhd}'
}} catch {{
    $_ | Out-File '{err}' -Encoding UTF8
    try {{ Dismount-VHD -Path '{vhd}' -ErrorAction SilentlyContinue }} catch {{}}
    throw
}}"#,
        vhd = vhd_abs.display(),
        out = output_dir_abs.display(),
        err = error_log_abs.display(),
    );

    std::fs::write(&script_path, &script)?;

    let status = std::process::Command::new("powershell")
        .args([
            "-Command",
            &format!(
                "Start-Process powershell -Verb RunAs -Wait -ArgumentList '-ExecutionPolicy', 'Bypass', '-File', '{}'",
                script_abs.display()
            ),
        ])
        .status()?;

    // Clean up the script
    std::fs::remove_file(&script_path).ok();

    if error_log.exists() {
        let error_text = std::fs::read_to_string(&error_log).unwrap_or_default();
        std::fs::remove_file(&error_log).ok();
        println!("WARNING: VHD extraction failed: {}", error_text.trim());
    } else if status.success() && output_dir.exists() {
        println!("Extracted to: {}", output_dir_abs.display());
        // Delete the VHD now that contents are extracted
        if let Err(e) = std::fs::remove_file(vhd_path) {
            println!("WARNING: Could not delete VHD: {e}");
        }
    } else {
        println!("WARNING: VHD extraction may have failed. Check UAC was accepted.");
    }

    Ok(())
}

/// Merge a differencing VHD into its parent using Hyper-V PowerShell cmdlets.
///
/// Steps:
/// 1. `Set-VHD` links the delta VHD to its parent
/// 2. `Merge-VHD` merges the delta's changes into the parent (modifies parent in-place)
///
/// Requires elevation (triggers UAC prompt).
#[cfg(windows)]
fn merge_vhd(base_vhd: &Path, delta_vhd: &Path) -> Result<()> {
    let base_abs = strip_extended_path_prefix(std::fs::canonicalize(base_vhd)?);
    let delta_abs = strip_extended_path_prefix(std::fs::canonicalize(delta_vhd)?);

    println!("Merging VHDs...");
    println!("  Base (parent): {}", base_abs.display());
    println!("  Delta (child): {}", delta_abs.display());

    // Set-VHD links the differencing disk to its parent, then
    // Merge-VHD (without -DestinationPath) merges the child into its immediate parent.
    let ps_commands = format!(
        "Set-VHD -Path '{}' -ParentPath '{}'; Merge-VHD -Path '{}' -Force",
        delta_abs.display(),
        base_abs.display(),
        delta_abs.display(),
    );

    let error_log = delta_vhd.with_extension("merge_error.txt");
    let error_log_abs = strip_extended_path_prefix(std::path::absolute(&error_log)?);

    // Wrap in try/catch to capture errors from the elevated process
    let wrapped = format!(
        "try {{ {} }} catch {{ $_ | Out-File '{}' -Encoding UTF8; throw }}",
        ps_commands,
        error_log_abs.display(),
    );

    let status = std::process::Command::new("powershell")
        .args([
            "-Command",
            &format!(
                "Start-Process powershell -Verb RunAs -Wait -ArgumentList '-Command', '{}'",
                wrapped.replace('\'', "''")
            ),
        ])
        .status()?;

    if error_log.exists() {
        let error_text = std::fs::read_to_string(&error_log).unwrap_or_default();
        std::fs::remove_file(&error_log).ok();
        println!("WARNING: VHD merge failed: {}", error_text.trim());
    } else if status.success() {
        println!("Merged delta into base VHD: {}", base_abs.display());
    } else {
        println!("WARNING: PowerShell exited with: {status}. Check UAC was accepted.");
    }

    Ok(())
}

#[derive(Parser)]
#[command(version, about = "decryptor for some SEGA containers", long_about = None)]
struct Cli {
    #[arg(long, help = "do not extract contents of decrypted image")]
    no_extract: bool,

    #[arg(required = true)]
    files: Vec<PathBuf>,
}

fn main() -> Result<()> {
    let cli = Cli::parse();

    // Pre-read bootids for all files to sort by sequence number (base first)
    let mut inputs: Vec<InputFile> = Vec::new();
    for path in &cli.files {
        let file = FscryptDecryptor::new(File::open(path)?).map_err(|e| anyhow!(e))?;
        inputs.push(InputFile {
            path: path.clone(),
            sequence_number: file.bootid.sequence_number,
        });
    }

    // Sort so seq=0 (base) is processed before seq>0 (deltas)
    inputs.sort_by_key(|f| f.sequence_number);

    // Track extracted VHDs for post-extraction merge
    let mut extracted_vhds: Vec<ExtractedVhd> = Vec::new();

    for input in &inputs {
        let path = &input.path;
        let file = FscryptDecryptor::new(File::open(path)?).map_err(|e| anyhow!(e))?;
        let bootid = file.bootid.clone();
        let output_filename = file.filename()?;
        let output_path = path.with_file_name(&output_filename);

        // Can't directly extract options since we can't impl exfat_fs::dir::ReadOffset
        // for our wrapper decryptor
        if cli.no_extract {
            let mut output_file = File::create(&output_path)?;

            output_file.set_len(file.len())?;

            let pb = ProgressBar::new(file.len())
                .with_style(
                    ProgressStyle::default_bar()
                        .template("{prefix} [{bar:20!.bright.yellow/dim.white}] {bytes:>8} [{elapsed}<{eta}, {bytes_per_sec}]")?
                );

            pb.set_prefix(output_filename.clone());

            let mut reader = BufReader::with_capacity(0x40000, file);

            loop {
                let buffer = reader.fill_buf()?;
                let length = buffer.len();

                if length == 0 {
                    break;
                }

                output_file.write_all(&buffer)?;
                reader.consume(length);

                pb.inc(length as u64);
            }

            output_file.flush()?;
        } else {
            match bootid.container_type {
                ContainerType::OS | ContainerType::APP => {
                    match extract_internal_vhd(&path, bootid.sequence_number) {
                        Ok(vhd_path) => {
                            let game_id =
                                std::str::from_utf8(&bootid.game_id)?.trim_end().to_string();
                            extracted_vhds.push(ExtractedVhd {
                                vhd_path,
                                input_path: path.clone(),
                                sequence_number: bootid.sequence_number,
                                game_id,
                            });
                        }
                        Err(e) => {
                            println!("WARNING: Failed to extract internal VHD: {e:#?}");
                        }
                    }
                }
                ContainerType::OPTION => match extract_exfat_contents(&path) {
                    Ok(_) => {}
                    Err(e) => {
                        println!("WARNING: Failed to extract exfat contents: {e:#?}");
                    }
                },
                _ => {
                    println!("WARNING: Unknown container type: {}", bootid.container_type);
                }
            }
        }
    }

    // Post-extraction: merge deltas and extract VHD contents (Windows only)
    #[cfg(windows)]
    if !cli.no_extract && !extracted_vhds.is_empty() {
        // Group by game_id
        let mut by_game: HashMap<String, Vec<&ExtractedVhd>> = HashMap::new();
        for vhd in &extracted_vhds {
            by_game.entry(vhd.game_id.clone()).or_default().push(vhd);
        }

        // Track which VHDs to extract contents from at the end
        let mut vhds_to_extract: Vec<PathBuf> = Vec::new();

        for (game_id, vhds) in &by_game {
            let base = vhds.iter().find(|v| v.sequence_number == 0);
            let deltas: Vec<_> = vhds.iter().filter(|v| v.sequence_number > 0).collect();

            if deltas.is_empty() {
                // Standalone base VHD, just extract its contents
                if let Some(base) = base {
                    vhds_to_extract.push(base.vhd_path.clone());
                }
                continue;
            }

            // If no base was extracted in this run, search the same directory
            let found_base;
            let base = match base {
                Some(b) => b,
                None => {
                    let delta_dir = deltas[0]
                        .input_path
                        .parent()
                        .unwrap_or_else(|| Path::new("."));
                    let pattern_prefix = format!("{game_id}_");

                    // First, look for an existing base VHD ({game_id}_*_0.vhd)
                    let existing_vhd = std::fs::read_dir(delta_dir)?
                        .filter_map(|e| e.ok())
                        .map(|e| e.path())
                        .find(|p| {
                            let name = p.file_name().unwrap_or_default().to_string_lossy();
                            name.starts_with(&pattern_prefix) && name.ends_with("_0.vhd")
                        });

                    if let Some(vhd_path) = existing_vhd {
                        println!("Found existing base VHD: {}", vhd_path.display());
                        found_base = ExtractedVhd {
                            vhd_path,
                            input_path: PathBuf::new(),
                            sequence_number: 0,
                            game_id: game_id.clone(),
                        };
                        &found_base
                    } else {
                        // Look for a base .app file ({game_id}_*_0.app) and extract it
                        let base_app = std::fs::read_dir(delta_dir)?
                            .filter_map(|e| e.ok())
                            .map(|e| e.path())
                            .find(|p| {
                                let name = p.file_name().unwrap_or_default().to_string_lossy();
                                name.starts_with(&pattern_prefix) && name.ends_with("_0.app")
                            });

                        match base_app {
                            Some(app_path) => {
                                println!(
                                    "Found base APP in same directory, extracting: {}",
                                    app_path.display()
                                );
                                let base_file = FscryptDecryptor::new(File::open(&app_path)?)
                                    .map_err(|e| anyhow!(e))?;
                                let base_seq = base_file.bootid.sequence_number;
                                drop(base_file);

                                match extract_internal_vhd(&app_path, base_seq) {
                                    Ok(vhd_path) => {
                                        found_base = ExtractedVhd {
                                            vhd_path,
                                            input_path: app_path,
                                            sequence_number: 0,
                                            game_id: game_id.clone(),
                                        };
                                        &found_base
                                    }
                                    Err(e) => {
                                        println!(
                                            "WARNING: Failed to extract base VHD: {e:#?}"
                                        );
                                        continue;
                                    }
                                }
                            }
                            None => {
                                println!(
                                    "WARNING: Delta VHD(s) found for {game_id} but no base (seq=0) found."
                                );
                                println!(
                                    "         Place the base .app or .vhd in the same directory."
                                );
                                continue;
                            }
                        }
                    }
                }
            };

            let mut merge_ok = true;
            for delta in deltas {
                if let Err(e) = merge_vhd(&base.vhd_path, &delta.vhd_path) {
                    println!(
                        "WARNING: Failed to merge {} with {}: {e:#?}",
                        delta.vhd_path.display(),
                        base.vhd_path.display()
                    );
                    merge_ok = false;
                }
            }

            // Extract the final merged VHD
            if merge_ok {
                vhds_to_extract.push(base.vhd_path.clone());
            }
        }

        // Extract all VHD contents to folders, then delete the VHDs
        for vhd_path in &vhds_to_extract {
            if vhd_path.exists() {
                if let Err(e) = extract_vhd_contents(vhd_path) {
                    println!("WARNING: Failed to extract VHD contents: {e:#?}");
                }
            }
        }
    }

    Ok(())
}