recovery_exporter.cpp

Go to the documentation of this file.

#include "recovery_exporter.hpp"
#include "disk_reader.hpp"
#include "file_carver.hpp"
#include "filesystem_recovery.hpp"
#include "rescue_imager.hpp"
#include "i18n.hpp"
 
#include <algorithm>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <utility>
#include <unordered_set>
 
namespace fs = std::filesystem;
 
namespace {
 
    bool export_range(DiskReader& reader,
                      uint64_t offset,
                      uint64_t size,
                      const fs::path& out_path,
                      std::string& error,
                      std::function<bool()> cancel_cb)
    {
        std::ofstream out(out_path, std::ios::binary);
        if (!out) {
            error = trs("Unable to create output file") + std::string(": ") + out_path.string();
            return false;
        }
 
        std::vector<unsigned char> buffer(1024 * 1024);
 
        uint64_t remaining = size;
        uint64_t cursor = offset;
 
        while (remaining > 0)
        {
            if (cancel_cb && cancel_cb())
            {
                error = trs("Recovery cancelled by user");
                return false;
            }
 
            const size_t chunk = static_cast<size_t>(
                std::min<uint64_t>(buffer.size(), remaining));
 
            if (!reader.read_at(cursor, buffer.data(), chunk))
            {
                error = trs("Read error while exporting partition image");
                return false;
            }
 
            out.write(reinterpret_cast<const char*>(buffer.data()),
                      static_cast<std::streamsize>(chunk));
 
            if (!out)
            {
                error = trs("Write error while exporting partition image");
                return false;
            }
 
            cursor += chunk;
            remaining -= chunk;
        }
 
        return true;
    }
 
    bool should_export(const ScanResult& r)
    {
        return r.kind == ResultKind::Partition ||
        r.kind == ResultKind::LostPartitionCandidate ||
        r.kind == ResultKind::FilesystemVolume ||
        r.kind == ResultKind::NtfsVolume;
    }
 
    std::string sanitize(const std::string& s)
    {
        std::string out;
 
        for (char c : s)
        {
            if (isalnum(static_cast<unsigned char>(c)))
                out.push_back(c);
            else if (c == ' ' || c == '-' || c == '_')
                out.push_back('_');
        }
 
        if (out.empty())
            out = "item";
 
        return out;
    }
 
    void write_recovery_manifest(
        const fs::path& output_dir,
        const ScanReport& report,
        const std::unordered_set<size_t>& selected_indices,
        bool filter_results,
        const RecoverySummary& summary,
        const std::vector<RecoveredFileRecord>& fs_recovered,
        const std::vector<CarvedFile>& carved,
        std::string& error)
    {
        std::ofstream out(output_dir / "recovery_manifest.txt", std::ios::binary);
 
        if (!out)
        {
            error = trs("Unable to create output file") + std::string(": ") +
            (output_dir / "recovery_manifest.txt").string();
            return;
        }
 
        out << "input=" << report.input_path << "\n";
        out << "selected_filter=" << (filter_results ? 1 : 0) << "\n";
        out << "selected_results=" << selected_indices.size() << "\n";
        out << "exported_images=" << summary.exported_images.size() << "\n";
        out << "recovered_files=" << summary.recovered_files.size() << "\n\n";
 
        out << "[filesystem_recovery]\n";
 
        for (const auto& rf : fs_recovered)
        {
            out << rf.filesystem << '\t'
            << rf.integrity << '\t'
            << rf.size_bytes << '\t'
            << rf.original_path << '\t'
            << rf.output_path;
 
            if (!rf.note.empty())
                out << '\t' << rf.note;
 
            out << '\n';
        }
 
        out << "\n[carved]\n";
 
        for (const auto& cf : carved)
        {
            out << cf.type << '\t'
            << cf.size << '\t'
            << cf.output_path << '\n';
        }
    }
 
    std::vector<std::pair<uint64_t, uint64_t>> selected_carve_ranges(
        const ScanReport& report,
        const std::unordered_set<size_t>& selected_indices,
        bool filter_results,
        uint64_t source_size)
    {
        std::vector<std::pair<uint64_t, uint64_t>> ranges;
 
        auto add_range = [&](uint64_t offset, uint64_t size)
        {
            if (size == 0 || offset >= source_size)
                return;
 
            size = std::min<uint64_t>(size, source_size - offset);
 
            if (ranges.empty())
            {
                ranges.emplace_back(offset, size);
                return;
            }
 
            auto& last = ranges.back();
 
            const uint64_t last_end = last.first + last.second;
 
            if (offset <= last_end)
            {
                const uint64_t new_end = std::max<uint64_t>(last_end, offset + size);
                last.second = new_end - last.first;
            }
            else
            {
                ranges.emplace_back(offset, size);
            }
        };
 
        if (!filter_results)
        {
            ranges.emplace_back(0, source_size);
            return ranges;
        }
 
        std::vector<std::pair<uint64_t, uint64_t>> temp;
 
        for (size_t idx = 0; idx < report.results.size(); ++idx)
        {
            if (!selected_indices.count(idx))
                continue;
 
            const auto& r = report.results[idx];
 
            if (r.offset_bytes == 0 && r.size_bytes == 0)
                continue;
 
            switch (r.kind)
            {
                case ResultKind::Partition:
                case ResultKind::LostPartitionCandidate:
                case ResultKind::FilesystemCandidate:
                case ResultKind::FilesystemVolume:
                case ResultKind::NtfsVolume:
                case ResultKind::NtfsMetadata:
                    temp.emplace_back(r.offset_bytes, r.size_bytes);
                    break;
                default:
                    break;
            }
        }
 
        if (temp.empty())
            return ranges;
 
        std::sort(temp.begin(), temp.end(),
                  [](const auto& a, const auto& b)
                  {
                      return a.first < b.first;
                  });
 
        for (const auto& r : temp)
            add_range(r.first, r.second);
 
        return ranges;
    }
 
} // namespace
 
 
 
bool RecoveryExporter::create_rescue_image(
    const std::string& input_path,
    const std::string& output_image_path,
    RecoverySummary& summary,
    std::vector<std::string>& logs,
    std::string& error,
    std::function<bool()> cancel_cb)
{
    RescueImageReport report;
    if (!RescueImager::create_image(input_path, output_image_path, report, error, {}, cancel_cb))
    {
        return false;
    }
 
    summary.rescue_image_path = report.image_path;
    summary.rescue_map_path = report.map_path;
    logs.insert(logs.end(), report.logs.begin(), report.logs.end());
    return true;
}
 
 
bool RecoveryExporter::export_results_and_carve(
    const ScanReport& report,
    const std::string& output_dir,
    const std::vector<size_t>& selected_result_indices,
    RecoverySummary& summary,
    std::vector<std::string>& logs,
    std::string& error,
    std::function<bool()> cancel_cb)
{
    summary = {};
 
    try
    {
        fs::create_directories(output_dir);
        fs::create_directories(fs::path(output_dir) / "images");
        fs::create_directories(fs::path(output_dir) / "carved");
    }
    catch (const std::exception& ex)
    {
        error = ex.what();
        return false;
    }
 
    if (cancel_cb && cancel_cb())
    {
        logs.push_back(trs("Recovery cancelled by user."));
        return false;
    }
 
    DiskReader reader;
 
    std::string open_error;
 
    if (!reader.open_readonly(report.input_path, open_error))
    {
        error = trs("Unable to open source");
 
        if (!open_error.empty())
            error += std::string(": ") + open_error;
 
        logs.push_back(trs("Recovery open error:") + std::string(" ") + error);
 
        return false;
    }
 
    std::unordered_set<size_t> selected_indices(
        selected_result_indices.begin(),
                                                selected_result_indices.end());
 
    const bool filter_results = !selected_indices.empty();
 
    size_t index = 0;
 
    for (size_t result_index = 0;
         result_index < report.results.size();
    ++result_index)
         {
             if (cancel_cb && cancel_cb())
             {
                 logs.push_back(trs("Recovery cancelled during image export."));
                 return false;
             }
 
             const auto& r = report.results[result_index];
 
             if (filter_results && !selected_indices.count(result_index))
                 continue;
 
             if (!should_export(r) || r.size_bytes == 0)
                 continue;
 
             if (r.offset_bytes + r.size_bytes > reader.size())
                 continue;
 
             std::ostringstream name;
 
             name << std::setw(3)
             << std::setfill('0')
             << index++
             << '_'
             << sanitize(r.type)
             << '_'
             << r.offset_bytes
             << ".img";
 
             fs::path out_path =
             fs::path(output_dir) / "images" / name.str();
 
             if (!export_range(
                 reader,
                 r.offset_bytes,
                 r.size_bytes,
                 out_path,
                 error,
                 cancel_cb))
             {
                 return false;
             }
 
             summary.exported_images.push_back(out_path.string());
 
             logs.push_back(
                 trs("Exported volume image:") +
                 std::string(" ") +
                 out_path.string());
         }
 
         if (cancel_cb && cancel_cb())
         {
             logs.push_back(trs("Recovery cancelled before filesystem recovery."));
             return false;
         }
 
         ScanReport filtered_report = report;
 
         if (filter_results)
         {
             filtered_report.results.clear();
 
             for (size_t idx = 0; idx < report.results.size(); ++idx)
             {
                 if (selected_indices.count(idx))
                     filtered_report.results.push_back(report.results[idx]);
             }
         }
 
         std::vector<RecoveredFileRecord> fs_recovered;
 
         if (!FilesystemRecovery::recover_supported_filesystems(
             reader,
             filtered_report,
             output_dir,
             fs_recovered,
             logs,
             error))
         {
             return false;
         }
 
         for (const auto& rf : fs_recovered)
             summary.recovered_files.push_back(rf.output_path);
 
    if (cancel_cb && cancel_cb())
    {
        logs.push_back(trs("Recovery cancelled during filesystem recovery."));
        return false;
    }
 
    const auto carve_ranges =
    selected_carve_ranges(
        report,
        selected_indices,
        filter_results,
        reader.size());
 
    std::vector<CarvedFile> carved;
 
    if (!carve_ranges.empty())
    {
        if (!FileCarver::carve_ranges_to_directory(
            reader,
            carve_ranges,
            (fs::path(output_dir) / "carved").string(),
                                                   carved,
                                                   logs,
                                                   error))
        {
            return false;
        }
    }
 
    for (const auto& cf : carved)
        summary.recovered_files.push_back(cf.output_path);
 
    if (cancel_cb && cancel_cb())
    {
        logs.push_back(trs("Recovery cancelled during carving."));
        return false;
    }
 
    write_recovery_manifest(
        fs::path(output_dir),
                            report,
                            selected_indices,
                            filter_results,
                            summary,
                            fs_recovered,
                            carved,
                            error);
 
    if (!error.empty())
        return false;
 
    logs.push_back(trs("Recovery export completed."));
 
    return true;
}