lowlevelscanner.cpp

#include "lowlevelscanner.h"
#include "bip39_sequence.h"
#include "bip39_checksum.h"
#include "bip39_wordlist_raw.h"
#include "bip39_wordlist_std.h"
#include <QFile>
#include <QTextStream>
#include <QDateTime>
#include <QDebug>
#include <QThread>

LowLevelScanner::LowLevelScanner(QObject *parent) : QObject(parent) {
    wordSet = Bip39Sequence::buildWordSet(bip39_wordlist_raw::wordlist);
}
LowLevelScanner::~LowLevelScanner() { closeDevice(); }

void LowLevelScanner::startScan(const QString &devicePath, const QString &outputDir, quint64 blockSize) {
    blockSizeBytes = blockSize;
    outputDirectory = outputDir;
    currentBlock = 0;
    isScanning = true;
    matches.clear();
    foundMatches.clear();

    qDebug() << "[LowLevelScanner] Opening device:" << devicePath;

    if (!openDevice(devicePath)) {
        qWarning() << "[LowLevelScanner] Failed to open device:" << devicePath;
        emit scanFinished();
        return;
    }

    // Determine total blocks
    LARGE_INTEGER size;
    size.QuadPart = 0;
    GET_LENGTH_INFORMATION lengthInfo;
    ZeroMemory(&lengthInfo, sizeof(lengthInfo));
    DWORD bytesReturned = 0;

    if (DeviceIoControl(deviceHandle, IOCTL_DISK_GET_LENGTH_INFO,
                        NULL, 0, &lengthInfo, sizeof(lengthInfo),
                        &bytesReturned, NULL)) {
        size.QuadPart = lengthInfo.Length.QuadPart;
        totalBlocks = size.QuadPart / blockSizeBytes;
        emit totalBlocksKnown(totalBlocks);
    } else {
        qWarning() << "[LowLevelScanner] DeviceIoControl failed, using fallback.";
        totalBlocks = 1024;
        emit totalBlocksKnown(totalBlocks);
    }

    scanLoop();
    closeDevice();
    emit scanFinished();
}

void LowLevelScanner::stopScan() {
    if (isScanning) {
        isScanning = false;
        // Wake the scan thread if it is waiting on a pause so it can exit cleanly.
        {
            QMutexLocker locker(&pauseMutex);
            isPaused = false;
            pauseCondition.wakeAll();
        }
        qDebug() << "[LowLevelScanner] stopScan() called, writing partial CSV...";
        writePartialCsv();
        emit scanAborted();
    }
}

void LowLevelScanner::pauseScan() {
    QMutexLocker locker(&pauseMutex);
    isPaused = true;
    qDebug() << "[LowLevelScanner] Scan paused at block" << currentBlock;
}

void LowLevelScanner::resumeScan() {
    QMutexLocker locker(&pauseMutex);
    isPaused = false;
    pauseCondition.wakeAll();
    qDebug() << "[LowLevelScanner] Scan resumed at block" << currentBlock;
}

bool LowLevelScanner::openDevice(const QString &devicePath) {
    closeDevice();
    deviceHandle = CreateFileW((LPCWSTR)devicePath.utf16(),
                               GENERIC_READ,
                               FILE_SHARE_READ | FILE_SHARE_WRITE,
                               NULL, OPEN_EXISTING, 0, NULL);
    return deviceHandle != INVALID_HANDLE_VALUE;
}

void LowLevelScanner::closeDevice() {
    if (deviceHandle != INVALID_HANDLE_VALUE) {
        CloseHandle(deviceHandle);
        deviceHandle = INVALID_HANDLE_VALUE;
    }
}

void LowLevelScanner::scanLoop() {
    QByteArray buffer;
    buffer.resize((int)blockSizeBytes);
    DWORD bytesRead = 0;

    qDebug() << "[LowLevelScanner] Starting scanLoop with block size:" << blockSizeBytes << "bytes";

    while (isScanning && currentBlock < totalBlocks) {
        {
            QMutexLocker locker(&pauseMutex);
            while (isPaused && isScanning)
                pauseCondition.wait(&pauseMutex);
        }
        if (!isScanning) break;

        LARGE_INTEGER offset;
        offset.QuadPart = currentBlock * blockSizeBytes;
        SetFilePointerEx(deviceHandle, offset, NULL, FILE_BEGIN);

        bool readOk = ReadFile(deviceHandle, buffer.data(), (DWORD)blockSizeBytes, &bytesRead, NULL);
        if (!readOk || bytesRead == 0) {
            qWarning() << "[LowLevelScanner] ReadFile failed at block" << currentBlock;
            break;
        }

        if (findBip39Words(buffer, offset.QuadPart)) {
            emit blockMatchFound(currentBlock);  // Highlight block in UI
        }

        emit progressUpdated(currentBlock, totalBlocks);
        currentBlock++;
    }

    qDebug() << "[LowLevelScanner] scanLoop finished.";
    if (!isScanning) {
        writePartialCsv();  // Save partial results if scan stopped early
    }
}

bool LowLevelScanner::findBip39Words(const QByteArray &data, quint64 offset) {
    // -----------------------------------------------------------------------
    // Phase 1 - word-pair detection (existing heuristic).
    //   Drives UI signals and determines whether to proceed to phase 2.
    //   No file is dumped here.
    // -----------------------------------------------------------------------
    bool pairFound = false;
    QStringList wordsFoundInBlock;
    const int proximity = 6;

    for (const QByteArray &word1 : bip39_wordlist_raw::wordlist) {
        int index1 = data.indexOf(word1);
        while (index1 != -1) {
            if (isValidWordMatch(data, index1, word1.size())) {
                QString word1Str = QString(word1);
                quint64 abs1 = offset + index1;

                int cs1 = qMax(0, index1 - 16);
                int cl1 = qMin(data.size() - cs1, index1 - cs1 + (int)word1.size() + 16);
                emit wordFoundInBlock(word1Str, currentBlock, abs1,
                                      data.mid(cs1, cl1).toHex(' ').toUpper());
                if (!wordsFoundInBlock.contains(word1Str))
                    wordsFoundInBlock.append(word1Str);

                int searchStart = index1 + word1.size();
                int searchEnd   = qMin(searchStart + proximity, data.size());
                QByteArray window = data.mid(searchStart, searchEnd - searchStart);

                for (const QByteArray &word2 : bip39_wordlist_raw::wordlist) {
                    int index2 = window.indexOf(word2);
                    if (index2 == -1) continue;
                    int abs2idx = searchStart + index2;
                    if (!isValidWordMatch(data, abs2idx, word2.size())) continue;

                    QString word2Str = QString(word2);
                    quint64 abs2 = offset + abs2idx;

                    int cs2 = qMax(0, abs2idx - 16);
                    int cl2 = qMin(data.size() - cs2, abs2idx - cs2 + (int)word2.size() + 16);
                    emit wordFoundInBlock(word2Str, currentBlock, abs2,
                                          data.mid(cs2, cl2).toHex(' ').toUpper());
                    if (!wordsFoundInBlock.contains(word2Str))
                        wordsFoundInBlock.append(word2Str);

                    qDebug() << "[LowLevelScanner] Word pair:" << word1 << word2
                             << "at block offset" << index1;
                    pairFound = true;
                    break;
                }
            }
            index1 = data.indexOf(word1, index1 + word1.size());
        }
    }

    if (!wordsFoundInBlock.isEmpty())
        emit blockMatchComplete(currentBlock, wordsFoundInBlock.size(), wordsFoundInBlock);

    if (!pairFound)
        return false;

    // Word pair confirmed - notify UI to highlight yellow (before checksum).
    emit blockPairFound(currentBlock);

    // -----------------------------------------------------------------------
    // Phase 2 - sequence extraction + BIP-39 checksum validation.
    //   Only blocks that already contain a word pair reach this point.
    //   A file is dumped only when a complete phrase passes checksum.
    // -----------------------------------------------------------------------
    const QVector<Bip39Sequence::Match> sequences = Bip39Sequence::extract(data, wordSet);
    for (const Bip39Sequence::Match &seq : sequences) {
        std::vector<std::string> stdWords;
        stdWords.reserve(seq.words.size());
        for (const QString &w : seq.words)
            stdWords.push_back(w.toStdString());

        if (!BIP39Checksum::validate(stdWords, Bip39WordsStd::wordlist))
            continue;

        // Valid phrase - dump the surrounding region and record the match.
        quint64 seqOffset = offset + (quint64)seq.byteOffset;
        qDebug() << "[LowLevelScanner] Checksum validated:" << seq.words.size()
                 << "words at absolute offset" << seqOffset;

        const int context = 25000;
        int start  = qMax(0, seq.byteOffset - context);
        int length = qMin(data.size() - start, 2 * context + 200);
        dumpMatch(data.mid(start, length), offset + start);

        foundMatches.push_back(
            QString("%1,%2").arg(seqOffset).arg(seq.words.join(' ')));
        matches.append(qMakePair(seq.words.join(' '), seqOffset));
        return true;   // one validated phrase per block is sufficient
    }

    // Pair found but no phrase survived checksum - suppress the dump.
    qDebug() << "[LowLevelScanner] Pair in block" << currentBlock
             << "failed checksum - suppressed.";
    return false;
}

bool LowLevelScanner::isValidWordMatch(const QByteArray &data, int index, int wordLength) {
    auto isDelimiter = [](char c) {
        return isspace(c) || ispunct(c) || c == '\0';
    };
    char before = (index > 0) ? data[index - 1] : ' ';
    char after = (index + wordLength < data.size()) ? data[index + wordLength] : ' ';
    return isDelimiter(before) && isDelimiter(after);
}

void LowLevelScanner::dumpMatch(const QByteArray &data, quint64 matchOffset) {
    QString filename = QString("%1/match_%2_offset_%3.txt")
    .arg(outputDirectory)
        .arg(QDateTime::currentDateTime().toString("yyyyMMdd_HHmmss"))
        .arg(matchOffset);

    QFile file(filename);
    if (file.open(QIODevice::WriteOnly)) {
        file.write(data);
        file.close();
        qDebug() << "[LowLevelScanner] Match dumped to" << filename << "with size" << data.size();
        emit matchFound(filename, matchOffset);
    }
}

void LowLevelScanner::writePartialCsv() {
    if (matches.isEmpty()) {
        qDebug() << "[LowLevelScanner] No matches to save in partial CSV.";
        return;
    }

    QString timestamp = QDateTime::currentDateTime().toString("yyyyMMdd_HHmmss");
    QString csvPath = QString("%1/matches_partial_%2.csv")
                          .arg(outputDirectory)
                          .arg(timestamp);

    QFile csvFile(csvPath);
    if (csvFile.open(QIODevice::WriteOnly | QIODevice::Text)) {
        QTextStream out(&csvFile);
        out << "Word,Offset\n";
        for (const auto &m : matches) {
            out << m.first << "," << m.second << "\n";
        }
        csvFile.close();
        qDebug() << "[LowLevelScanner] Partial CSV saved:" << csvPath;
        emit partialCsvSaved(csvPath);
    } else {
        qWarning() << "[LowLevelScanner] Failed to save partial CSV at:" << csvPath;
    }
}