#include "../core/module.h"
#include "../core/ui.h"
#include <HardwareSerial.h>

class ExploitChain : public Module {
  enum Stage { ARMED, RUN_CMD1, RUN_CMD2, RUN_CMD3, RUN_CMD4, DONE } stage = ARMED;
  bool running = false;
  uint32_t stageTime = 0;
  uint32_t txCount = 0;
  char responses[3][40] = {{0},{0},{0}};
  HardwareSerial ser;
  uint32_t rxBytes = 0;

public:
  const char* name() const override { return "Exploit Chain"; }
  const char* blurb() const override { return "manual exploit send"; }

  void onEnter() override {
    running = false;
    stage = ARMED;
    stageTime = 0;
    txCount = 0;
    rxBytes = 0;
    memset(responses, 0, sizeof(responses));
    say("UART ready: [space] to send payloads");
    ser.begin(115200, SERIAL_8N1, 16, 17);
  }
  void onExit() override { running = false; ser.end(); }

  bool onKey(char c) override {
    if (c == ' ') { running = !running; if (running) { stage = RUN_CMD1; stageTime = 0; } return true; }
    if (c == 'r') { say("Reset"); stage = ARMED; running = false; stageTime = 0; return true; }
    return false;
  }

  void tick() override {
    if (!running) return;

    stageTime++;

    // Read any incoming data
    while (ser.available() && rxBytes < 2000) rxBytes += ser.read();

    switch (stage) {
      case RUN_CMD1:
        if (stageTime == 5) { ser.println("id"); txCount++; say("TX: id"); }
        if (stageTime == 50) { stage = RUN_CMD2; stageTime = 0; }
        break;
      case RUN_CMD2:
        if (stageTime == 5) { ser.println("uname -a"); txCount++; say("TX: uname -a"); }
        if (stageTime == 50) { stage = RUN_CMD3; stageTime = 0; }
        break;
      case RUN_CMD3:
        if (stageTime == 5) { ser.println("cat /proc/version"); txCount++; say("TX: cat /proc/version"); }
        if (stageTime == 50) { stage = RUN_CMD4; stageTime = 0; }
        break;
      case RUN_CMD4:
        if (stageTime == 5) { ser.println("whoami"); txCount++; say("TX: whoami"); }
        if (stageTime == 50) { stage = DONE; running = false; say("Done: %u cmds, %lu rx", (unsigned)txCount, rxBytes); }
        break;
      case DONE: running = false; break;
      default: break;
    }
        break;
      case DONE:
        running = false;
        break;
    }
  }

  void draw() override {
    const char* sn[] = {"ARMED", "CMD1", "CMD2", "CMD3", "CMD4", "DONE"};
    ui::lineC(0, ui::accent(), "Exploit: %s  %s", sn[stage], running ? "ACTIVE" : "idle");
    ui::line(1, "sent: %u cmds  rx: %lu bytes", (unsigned)txCount, rxBytes);
    if (running) ui::bar(2, stageTime / 50.0f, ui::glow(), "progress");
    for (int i = 0; i < 3; i++) ui::line(4 + i, "%s", responses[i]);
    if (running) ui::spinner(228, ui::BODY_Y + 1, ui::glow());
    ui::hintBar("[space]execute [r]eset  [`]back");
  }

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

Module* makeExploitChain() { return new ExploitChain(); }