#include "unlock.h"
#include "scsi_structs.h"
#include "rc4.h"
#include "sha1.h"
#include "utils.h"
#include "xbox_ref_log.h"

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <windows.h>
#include <stddef.h>

#define printf xbox_ref_printf

static uint8_t scsibuffer[2001];
static TRC4Context rc4key;

// Helper function to log hex dumps cleanly to the console and to the untruncated log file.
static void LogHexDump(const char* prefix, const uint8_t* data, DWORD length) {
    size_t console_limit = 0;
    if (!data || length == 0) return;
    if (prefix && strcmp(prefix, "Received Data [Step 3: Read DVD Struct 0xAD]") == 0)
        console_limit = 32;
    xbox_ref_log_hexdump(prefix, data, (size_t)length, console_limit);
}

void UnlockDrive(HANDLE AHandle) {
    printf("Structure size: %zu\n", sizeof(SCSI_PASS_THROUGH_DIRECT));
    SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER SPTDW;
    DWORD Size = sizeof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER);
    DWORD Returned;
    BOOL Status;
    int i, k, l;
    int chalpos[11];
    uint8_t restable[261];
    uint8_t hash[0x2C];
    uint8_t shadigest[20];
    uint8_t shastring[20];

    // --- STEP 1: INITIAL READ CAPACITY ---
    ZeroMemory(&SPTDW, sizeof(SPTDW));
    ZeroMemory(scsibuffer, sizeof(scsibuffer));
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.SenseInfoLength = 32;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_IN;
    SPTDW.Spt.DataTransferLength = 8; 
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x25;

    LogHexDump("Sending CDB [Step 1: Init Read Cap]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
    
    if (Status) {
        LogHexDump("Received Data [Step 1: Init Read Cap]", scsibuffer, SPTDW.Spt.DataTransferLength);
        unsigned int maxLBA = (scsibuffer[0] << 24) |
                              (scsibuffer[1] << 16) | 
                              (scsibuffer[2] << 8)  | 
                               scsibuffer[3];
        printf("Initial Max LBA: %u\n", maxLBA);
        PrintFormattedCapacity(scsibuffer);
    } else {
        printf("Fatal error: Initial capacity read failed (Error %lu).\n", GetLastError());
        return;
    }

    // --- STEP 2: MODE SENSE (Page 0x3E) ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_IN;
    SPTDW.Spt.DataTransferLength = 28;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x5A;
    SPTDW.Spt.Cdb[2] = 0x3E;
    SPTDW.Spt.Cdb[8] = 0x1C;
    
    LogHexDump("Sending CDB [Step 2: Mode Sense]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
    if (Status) LogHexDump("Received Data [Step 2: Mode Sense]", scsibuffer, SPTDW.Spt.DataTransferLength);

    // --- STEP 3: READ DVD STRUCTURE (The Challenge Table) ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 12;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_IN;
    SPTDW.Spt.DataTransferLength = 0x664;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0xAD;
    SPTDW.Spt.Cdb[2] = 0xFF; SPTDW.Spt.Cdb[3] = 0x02; SPTDW.Spt.Cdb[4] = 0xFD;
    SPTDW.Spt.Cdb[5] = 0xFF; SPTDW.Spt.Cdb[6] = 0xFE; SPTDW.Spt.Cdb[8] = 0x06;
    SPTDW.Spt.Cdb[9] = 0x64; SPTDW.Spt.Cdb[11] = 0xC0;
    
    LogHexDump("Sending CDB [Step 3: Read DVD Struct 0xAD]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
    if (Status) LogHexDump("Received Data [Step 3: Read DVD Struct 0xAD]", scsibuffer, SPTDW.Spt.DataTransferLength);

    // Fallback Logic for Hitachi 4241N Drives
    if (!Status || scsibuffer[772] != 1 || scsibuffer[773] == 0) {
        printf("8050L Table Invalid. Attempting Hitachi 0xFD Fallback...\n");
        memset(&SPTDW, 0, sizeof(SPTDW));
        memset(scsibuffer, 0, 2000);
        SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
        SPTDW.Spt.CdbLength = 12;
        SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_IN;
        SPTDW.Spt.DataTransferLength = 0x664;
        SPTDW.Spt.TimeOutValue = 120;
        SPTDW.Spt.DataBuffer = scsibuffer;
        SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
        
        SPTDW.Spt.Cdb[0] = 0xFD; 
        SPTDW.Spt.Cdb[1] = 0x01; 
        SPTDW.Spt.Cdb[8] = 0x06;
        SPTDW.Spt.Cdb[9] = 0x64;

        LogHexDump("Sending CDB [Step 3: Read DVD Struct Fallback 0xFD]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
        Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
        if (Status) LogHexDump("Received Data [Step 3: Read DVD Struct Fallback 0xFD]", scsibuffer, SPTDW.Spt.DataTransferLength);

        if (!Status || scsibuffer[772] != 1) {
            printf("Fatal error: Invalid host challenge table (Both methods failed).\n");
            return;
        }
        printf("Success: Acquired Hitachi 0xFD Handshake Table!\n");
    }

    // --- CRYPTO: SHA1 & RC4 ---
    for (i = 0; i < 0x2C; i++) hash[i] = scsibuffer[0x4A3 + i];
    {
        SHA1_CTX ctx;
        SHA1_Init(&ctx);
        SHA1_Update(&ctx, hash, 0x2C);
        SHA1_Final(shadigest, &ctx);
    }
    memcpy(shastring, shadigest, 7);
    RC4Init(&rc4key, shastring, 7);
    for (i = 0; i <= 260; i++) restable[i] = scsibuffer[774 + i];
    rc4Decrypt(restable, restable, 0xFD, &rc4key);

    k = 0;
    for (l = 0; l <= 23; l++) {
        if (restable[l * 11] == 1) {
            chalpos[k] = l;
            k++;
        }
    }

    if (k < 2) {
        printf("Fatal error: Not enough usable challenge entries.\n");
        return;
    }

    // --- STEP 4: MODE SELECT (First Challenge) ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_OUT;
    SPTDW.Spt.DataTransferLength = 28;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x55; SPTDW.Spt.Cdb[8] = 0x1C;
    scsibuffer[1] = 0x1A; scsibuffer[8] = 0x3E;
    scsibuffer[9] = 0x12;
    scsibuffer[11] = 0x01; scsibuffer[13] = 0xD1; scsibuffer[14] = 0x01;
    memcpy(&scsibuffer[15], &restable[1 + chalpos[k - 2] * 11], 5);
    
    LogHexDump("Sending CDB [Step 4: Mode Select Chal 1]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    LogHexDump("Sending Data [Step 4: Mode Select Chal 1]", scsibuffer, SPTDW.Spt.DataTransferLength);
    DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);

    // --- STEP 5: MODE SENSE (Verify First Challenge) ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_IN;
    SPTDW.Spt.DataTransferLength = 28;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x5A; SPTDW.Spt.Cdb[2] = 0x3E; SPTDW.Spt.Cdb[8] = 0x1C;
    
    LogHexDump("Sending CDB [Step 5: Mode Sense Verify 1]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
    if (Status) LogHexDump("Received Data [Step 5: Mode Sense Verify 1]", scsibuffer, SPTDW.Spt.DataTransferLength);
    
    // --- STEP 6: MODE SELECT (Second Challenge) ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_OUT;
    SPTDW.Spt.DataTransferLength = 28;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x55; SPTDW.Spt.Cdb[8] = 0x1C;
    scsibuffer[1] = 0x1A; scsibuffer[8] = 0x3E;
    scsibuffer[9] = 0x12;
    scsibuffer[12] = 0x01; // Marker for second challenge
    memcpy(&scsibuffer[15], &restable[1 + chalpos[k - 1] * 11], 5);
    
    LogHexDump("Sending CDB [Step 6: Mode Select Chal 2]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    LogHexDump("Sending Data [Step 6: Mode Select Chal 2]", scsibuffer, SPTDW.Spt.DataTransferLength);
    DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);

    // --- STEP 7: MODE SENSE (Verify Second Challenge) ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_IN;
    SPTDW.Spt.DataTransferLength = 28;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x5A; SPTDW.Spt.Cdb[2] = 0x3E; SPTDW.Spt.Cdb[8] = 0x1C;
    
    LogHexDump("Sending CDB [Step 7: Mode Sense Verify 2]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
    if (Status) LogHexDump("Received Data [Step 7: Mode Sense Verify 2]", scsibuffer, SPTDW.Spt.DataTransferLength);

    // --- STEP 8: FINAL MODE SELECT (Partition 1 Unlock) ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_OUT;
    SPTDW.Spt.DataTransferLength = 28;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x55; SPTDW.Spt.Cdb[8] = 0x1C;
    scsibuffer[1] = 0x1A; scsibuffer[8] = 0x3E;
    scsibuffer[9] = 0x12;
    scsibuffer[10] = 0x01; scsibuffer[11] = 0x01; scsibuffer[12] = 0x01;
    scsibuffer[13] = 0xD1; scsibuffer[14] = 0x01;
    memcpy(&scsibuffer[15], &restable[1 + chalpos[k - 1] * 11], 5);
    
    LogHexDump("Sending CDB [Step 8: Mode Select Part1 Unlock]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    LogHexDump("Sending Data [Step 8: Mode Select Part1 Unlock]", scsibuffer, SPTDW.Spt.DataTransferLength);
    
    Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
    if (Status) {
        printf("Partition 1 Unlocked Successfully.\n");
    }

    // --- STEP 9: ENABLE STICKY DESCRAMBLING (Mode 0x31) ---
    printf("Enabling Sticky Descrambling (Mode 0x31)...\n");
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    scsibuffer[4] = 0x31; scsibuffer[5] = 0x06; scsibuffer[6] = 0x01;
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 6;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_OUT;
    SPTDW.Spt.DataTransferLength = 12;
    SPTDW.Spt.TimeOutValue = 10;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x15; 
    SPTDW.Spt.Cdb[1] = 0x11; 
    SPTDW.Spt.Cdb[4] = 12;
    
    LogHexDump("Sending CDB [Step 9: Sticky Descrambling]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    LogHexDump("Sending Data [Step 9: Sticky Descrambling]", scsibuffer, SPTDW.Spt.DataTransferLength);
    DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);

    // --- STEP 10: FINAL CAPACITY VERIFICATION ---
    memset(&SPTDW, 0, sizeof(SPTDW));
    memset(scsibuffer, 0, 2000);
    SPTDW.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    SPTDW.Spt.CdbLength = 10;
    SPTDW.Spt.DataIn = SCSI_IOCTL_DATA_IN;
    SPTDW.Spt.DataTransferLength = 8;
    SPTDW.Spt.TimeOutValue = 120;
    SPTDW.Spt.DataBuffer = scsibuffer;
    SPTDW.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);
    SPTDW.Spt.Cdb[0] = 0x25;
    
    LogHexDump("Sending CDB [Step 10: Final Verify]", SPTDW.Spt.Cdb, SPTDW.Spt.CdbLength);
    Status = DeviceIoControl(AHandle, IOCTL_SCSI_PASS_THROUGH_DIRECT, &SPTDW, Size, &SPTDW, Size, &Returned, NULL);
    
    if (Status) {
        LogHexDump("Received Data [Step 10: Final Verify]", scsibuffer, SPTDW.Spt.DataTransferLength);
        PrintFormattedCapacity(scsibuffer);
    }
}

// AI CODE
BOOL AuthenticateXboxMedia(HANDLE hDevice) {
    SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER sptdw = {0};
    DWORD bytesReturned;
    
    // We need a small buffer because 0xAD usually expects to return data, 
    // even if our custom handler just uses it as a trigger.
    uint8_t dummy_buffer[8]; 

    sptdw.Spt.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    sptdw.Spt.CdbLength = 12; // 0xAD is a 12-byte command
    sptdw.Spt.DataIn = SCSI_IOCTL_DATA_IN;
    sptdw.Spt.TimeOutValue = 10;
    sptdw.Spt.DataTransferLength = sizeof(dummy_buffer);
    sptdw.Spt.DataBuffer = dummy_buffer;
    sptdw.Spt.SenseInfoOffset = offsetof(SCSI_PASS_THROUGH_DIRECT_WITH_BUFFER, SenseBuf);

    // --- The 8050L / 4241N Hook Command ---
    sptdw.Spt.Cdb[0] = 0xAD; // READ DVD STRUCTURE
    sptdw.Spt.Cdb[2] = 0xFF; 
    sptdw.Spt.Cdb[3] = 0x02; 
    sptdw.Spt.Cdb[4] = 0xFD; // MAGIC BYTE: Our Gatekeeper at 0x90002670 listens for this
    sptdw.Spt.Cdb[11] = 0xC0;

    return DeviceIoControl(hDevice, IOCTL_SCSI_PASS_THROUGH_DIRECT, &sptdw, sizeof(sptdw), &sptdw, sizeof(sptdw), &bytesReturned, NULL);
}

// ORIGINAL CODE
BOOL ORIGINAL_AuthenticateXboxMedia(HANDLE hDevice) {
    SCSI_PASS_THROUGH_DIRECT sptd = {0};
    DWORD bytesReturned;
    unsigned char cdb[10] = {0};

    sptd.Length = sizeof(SCSI_PASS_THROUGH_DIRECT);
    sptd.CdbLength = 10;
    sptd.DataIn = SCSI_IOCTL_DATA_IN;
    sptd.TimeOutValue = 5;
    sptd.DataTransferLength = 0; // No data transfer for this command
    sptd.DataBuffer = NULL;

    // THE MAGIC LG GDR-8163B UNLOCK COMMAND
    sptd.Cdb[0] = 0xFF; // Vendor Specific
    sptd.Cdb[1] = 0x08; // Sub-command: Set Xbox Mode
    sptd.Cdb[2] = 0x01; // Enable
    // Note: Some firmwares use 0x01, some use 0xFF. Try 0x01 first.

    return DeviceIoControl(hDevice, IOCTL_SCSI_PASS_THROUGH_DIRECT, &sptd, sizeof(sptd), &sptd, sizeof(sptd), &bytesReturned, NULL);
}


BOOL KickXboxMediaAuth(HANDLE hDevice)
{
    if (AuthenticateXboxMedia(hDevice))
        return TRUE;

    return ORIGINAL_AuthenticateXboxMedia(hDevice);
}