#include "../core/module.h"
#include "../core/ui.h"

// DMA Attack Simulator: memory-to-memory transfers with privilege bypass.
// On systems with a DMA controller or I/O-MMU, attempt to:
// - Read/write arbitrary addresses
// - Bypass MPU/paging restrictions
// - Exfiltrate kernel memory
// - Inject code via DMA into code regions

class DmaAttack : public Module {
  enum Target { KERNEL_MEM, IOCTL_ARGS, PAGE_TABLE } target = KERNEL_MEM;
  uint32_t srcAddr = 0x40000000;   // Assume kernel region start
  uint32_t dstAddr = 0x20000000;   // User SRAM
  uint32_t size = 256;
  uint32_t transferred = 0;
  bool active = false;
  char msg[3][40] = {{0},{0},{0}};

public:
  const char* name() const override { return "DMA Attack"; }
  const char* blurb() const override { return "memory-to-memory with privesc"; }

  void onEnter() override {
    active = false;
    transferred = 0;
    say("DMA controller: probing...");
  }
  void onExit() override { active = false; }

  bool onKey(char c) override {
    if (c == 't') { target = (Target)((target + 1) % 3); return true; }
    if (c == '+') { size = (size * 2 > 4096) ? 256 : size * 2; return true; }
    if (c == 's') { startTransfer(); return true; }
    return false;
  }

  void tick() override {
    if (!active) return;
    if (transferred >= size) { active = false; say("-- transfer done --"); return; }

    // Simulate DMA: read from srcAddr, write to dstAddr
    // Bypass normal CPU cache/MMU on each chunk
    uint32_t chunk = 64;
    if (transferred + chunk > size) chunk = size - transferred;

    // Attempt unprotected read/write
    uint8_t* src = (uint8_t*)srcAddr;
    uint8_t* dst = (uint8_t*)dstAddr;

    // Disable cache during "transfer" (hardware normally does this)
    // memcpy(dst, src, chunk);  // Stub: real DMA would bypass MMU

    transferred += chunk;
  }

  void draw() override {
    const char* tn[] = {"KERNEL", "IOCTL", "PGTBL"};
    ui::lineC(0, ui::accent(), "DMA: %s  %s", tn[target], active ? "XFER" : "idle");
    ui::line(1, "src:0x%08lx dst:0x%08lx sz:%lu", (unsigned long)srcAddr,
             (unsigned long)dstAddr, (unsigned long)size);
    ui::bar(2, transferred / (float)size, ui::glow(), "dma");
    for (int i = 0; i < 3; i++) ui::line(4 + i, "%s", msg[i]);
    ui::hintBar("[t]arget [+]size [s]tart  [`]back");
  }

private:
  void say(const char* fmt, ...) {
    for (int i = 2; i > 0; i--) strncpy(msg[i], msg[i-1], 39);
    va_list ap; va_start(ap, fmt); vsnprintf(msg[0], 40, fmt, ap); va_end(ap);
  }

  void startTransfer() {
    // Attempt to configure DMA without privilege
    // Real systems use MMIO to program DMA, check:
    // - is DMA controller accessible from user space?
    // - are address restrictions enforced by I/O-MMU?

    transferred = 0;
    active = true;

    switch (target) {
      case KERNEL_MEM:
        srcAddr = 0x40000000;
        say("DMA: read kernel @0x%08lx", (unsigned long)srcAddr);
        break;
      case IOCTL_ARGS:
        srcAddr = 0x20010000;
        say("DMA: snoop IOCTL args");
        break;
      case PAGE_TABLE:
        srcAddr = 0xC0000000;  // Assume kernel page table
        say("DMA: exfil page table");
        break;
    }
  }
};

Module* makeDma() { return new DmaAttack(); }