// pico2 OTP glitcher firmware
// This firmware is designed to control a GPIO pin to create a glitch (short drop) for
// a specified duration when triggered via USB serial commands. It also includes a timeout
// mechanism to automatically turn off the target after a specified duration.
//
// To compile use pico-sdk and the CMakefiles.txt of Pico2 documentation.
//
// From https://jemos.net/2026/02/rp2350_otp_attack_experiment/pico-otp-glitcher.html .

#include <stdio.h>
#include "pico/stdlib.h"
#include "hardware/gpio.h"
#include "hardware/timer.h"


#define DEFAULT_RESET_GPIO 2
#define DEFAULT_POWER_GPIO 4
#define TRIGGER_CMD "t"
// Drop duration in microseconds
#define DROP_US 100
#define DEFAULT_TIMEOUT_US 5000000 // Default timeout of 5 seconds in microseconds

static uint reset_gpio = DEFAULT_RESET_GPIO;

static bool reset_active = true;
static alarm_id_t timeout_alarm_id = 0;

void drop_gpio(void) {
		// Set output latch low first, then enable output. This avoids a brief
		// unintended high level when switching from input (high-Z) to output.
		gpio_put(reset_gpio, 0);
		gpio_set_dir(reset_gpio, GPIO_OUT);
		sleep_us(DROP_US);
		// Return to high-impedance (input) so the pin is idle and not driven
		gpio_set_dir(reset_gpio, GPIO_IN);
}

void set_target_off(void) {
    gpio_put(reset_gpio, 0);
    gpio_set_dir(reset_gpio, GPIO_OUT);
    reset_active = true;
}

void set_target_on(void) {
	// Set pin in high-impedance, which by default is pulled-high (reset disabled)
	gpio_set_dir(reset_gpio, GPIO_IN);
    reset_active = false;
}

// Returns reset pin status which should reflect how target is.
#define TARGET_IS_ON false
#define TARGET_IS_OFF true
bool get_target_status(void) {
	return (gpio_get_dir(DEFAULT_POWER_GPIO) == GPIO_IN);
}

int64_t timeout_callback(alarm_id_t id, void *user_data) {
    if ( get_target_status() == TARGET_IS_ON) {
        set_target_off();
        printf("[glitcher] Timeout expired, turned off the target.\n");
    }
    return 0; // Do not reschedule
}

void handle_on_command(char *buf)
{
    unsigned int timeout = DEFAULT_TIMEOUT_US;
    if (sscanf(buf + 2, "%u", &timeout) == 1) {
        printf("[glitcher] Timeout set to %u microseconds.\n", timeout);
    } else {
        printf("[glitcher] Using default timeout of %u microseconds.\n", DEFAULT_TIMEOUT_US);
    }

    set_target_on();

    // Cancel any existing alarm
    if (timeout_alarm_id > 0) {
        cancel_alarm(timeout_alarm_id);
    }

    // Set a new alarm for the timeout
    timeout_alarm_id = add_alarm_in_us(timeout, timeout_callback, NULL, false);
	printf("[glitcher] Alarm set for %u microseconds.\n", timeout);
}

void handle_off_command(void) {
    if (timeout_alarm_id > 0) {
        cancel_alarm(timeout_alarm_id);
        timeout_alarm_id = 0;
    }
    set_target_off();
    printf("[glitcher] Target manually turned off.\n");
}

void show_prompt(void) {
	printf("\n[pico-otp-glitcher] Ready.\n");
	printf("Commands:\n");
	printf("  s <gpio>   Set reset GPIO (default: %u)\n", DEFAULT_RESET_GPIO);
	printf("  t          Trigger reset pulse\n");
	printf("  h          Show this help\n");
}

int main() {
	stdio_init_all();
	// Wait for USB connection
	int usb_connected = 0;
	for (int i = 0; i < 1000; ++i) {
		if (stdio_usb_connected()) {
			usb_connected = 1;
			break;
		}
		sleep_ms(10);
	}
	if (!usb_connected) {
		printf("[glitcher] Warning: USB not detected, continuing anyway...\n");
	}

	// Initialize the gpio but leave it as input (high-Z) so it is not driven when idle
	gpio_init(reset_gpio);
	gpio_set_dir(reset_gpio, GPIO_IN);

	// Initialize the GPIO used as power
	gpio_init(DEFAULT_POWER_GPIO);
	// At init it's an input (high-Z) — read it first to see external state
	int pre_level = gpio_get(DEFAULT_POWER_GPIO);
	printf("[glitcher] power GPIO %u pre-init level=%d\n", DEFAULT_POWER_GPIO, pre_level);

	gpio_set_dir(DEFAULT_POWER_GPIO, GPIO_OUT);
	gpio_put(DEFAULT_POWER_GPIO, 1);
	// Small delay to let the output settle on the pin and external circuits
	sleep_ms(1);
	int post_level = gpio_get(DEFAULT_POWER_GPIO);
	printf("[glitcher] power GPIO %u after set HIGH level=%d\n", DEFAULT_POWER_GPIO, post_level);
	if (!post_level) {
		printf("[glitcher] Warning: GPIO %u did not read back high after gpio_put(1).\n", DEFAULT_POWER_GPIO);
		printf("[glitcher] Possible causes: external pull-down, short to GND, or pin reserved by on-board peripheral.\n");
	}

	show_prompt();
	printf("> ");

	char buf[32];
	while (true) {
		
		int idx = 0;
		int cto;
        char c;

		// Read line from serial
		while (idx < (int)(sizeof(buf) - 1)) {
			cto = getchar_timeout_us(0);
			if (cto == PICO_ERROR_TIMEOUT) {
				tight_loop_contents();
				continue;
			}
			c = (char)cto;
			if (c == '\r' || c == '\n') break;

			// Handle backspace (ASCII 8 or 127)
			if ((c == 8 || c == 127)) {
				if (idx > 0) {
					idx--;
					// Move cursor back, erase char, move back again
					printf("\b \b");
				}
				continue;
			}

			buf[idx++] = c;
			putchar_raw(c);
		}
		printf("\n");

		// Clean trailing newline or carriage return
		if (idx > 0 && (buf[idx-1] == '\r' || buf[idx-1] == '\n')) {
			buf[idx-1] = '\0';
		} else {
			buf[idx] = '\0';
		}

		// Handle commands
		if (idx == 0 || (idx == 1 && buf[0] == 'h')) {
			show_prompt();
			printf("> ");
			continue;
		}

		if ( buf[0] == 't' ) {
			printf("[glitcher] Trigger received, dropping GPIO %u...\n", reset_gpio);
			drop_gpio();
			printf("[glitcher] Done.\n");
			printf("> ");
		} else if (buf[0] == 's' && (buf[1] == ' ' || buf[1] == '\t')) {
			int new_gpio = DEFAULT_RESET_GPIO;
			if (sscanf(buf + 2, "%d", &new_gpio) == 1 && new_gpio >= 0 && new_gpio <= 29) {
				// Initialize but keep the new gpio in high-Z (input) while idle.
				// It will be actively driven low only when a trigger occurs.
				gpio_init(new_gpio);
				gpio_set_dir(new_gpio, GPIO_IN);
				reset_gpio = (uint)new_gpio;
				printf("[glitcher] Reset GPIO set to %u\n", reset_gpio);
			} else {
				printf("[glitcher] Invalid GPIO.\n");
			}
			printf("> ");
		} else if ( buf[0] == 'o' ) {
    		handle_on_command(buf);
			printf("> ");
		} else if ( buf[0] == 'f' ) {
    		handle_off_command();
			printf("> ");
		} else {
			printf("[glitcher] Unknown command. Type 'h' for help.\n");
			printf("> ");
		}
	}
	return 0;
}