Motenpoche/msc_content/mep.c

#include <stdint.h>
#include <wolfssl/options.h>
#include <wolfssl/wolfcrypt/settings.h>
#include <wolfssl/wolfcrypt/ecc.h>
#include <wolfssl/wolfcrypt/pwdbased.h>
#include <wolfssl/wolfcrypt/sha512.h>
#include <wolfssl/wolfcrypt/chacha.h>

#include "../src/cryptoengine.h"

#include <fcntl.h>
#include <err.h>
#include <termio.h>
#include <linux/serial.h>
#include <stdio.h>
#include <errno.h>
#include <sys/poll.h>

#define HOMEPATH_PREFIX "~/.pvault"
//#define HOMEPATH_PREFIX "/root/.pvault"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MAX_SVC_NAME_LEN 16

static char pass_prompt[200] = "";
static char pass_confirm[200] = "";
static const char pass_prompt_phrase[] = "Enter your passphrase: ";
static const char pass_confirm_phrase[] = "Confirm your passphrase : ";

static const char pass_prompt_svc[] = "Password: ";
static const char pass_confirm_svc[] = "Confirm Password: ";

static int sfd = -1;;
static int pk_fd = -1;

static struct vault_status Status = { };
static int get_vault_status(void);


struct __attribute__((packed)) tofu_packet {
    uint32_t cdc_magic; /* 0x5afeca5e */
    uint16_t cdc_cmd;
    uint16_t cdc_len;
    struct vault_header vh;
};

void clearScreen() {
    system("clear");
}

void disableEcho() {
    struct termios term;
    tcgetattr(0, &term);
    term.c_lflag &= ~ECHO;
    tcsetattr(0, TCSANOW, &term);
}

void enableEcho() {
    struct termios term;
    tcgetattr(0, &term);
    term.c_lflag |= ECHO;
    tcsetattr(0, TCSANOW, &term);
}

void printPrompt() {
    printf(pass_prompt);
}

static void getPassword(char *password, int size) {
    int i = 0;
    char c;

    while (i < size - 1) {
        c = getchar();
        if (c == '\n') {
            break;
        }
        password[i++] = c;
        if (i >= 99)
            break;
    }

    password[i] = '\0';
}

int confirmPassphrase(const char *passphrase) {
    char confirmation[100];

    printf(pass_confirm);
    getPassword(confirmation, sizeof(confirmation));

    return strcmp(passphrase, confirmation) == 0;
}

static char passphrase[100];
static int askpass(void)
{
    int ret = 0;
    disableEcho();
    printPrompt();
    getPassword(passphrase, sizeof(passphrase));
    printf("\n");

    if (confirmPassphrase(passphrase)) {
        printf("\nConfirmed.\n");
        ret = 1;
    } else {
        printf("These passwords don't match. Please try again.\n");
    }
    enableEcho();
    return ret;
}



static int rate_to_constant(int baudrate) {
#ifdef __MACH__
#define B(x) case x: return x
#else
#define B(x) case x: return B##x
#endif
    switch(baudrate) {
        B(50);     B(75);     B(110);    B(134);    B(150);
        B(200);    B(300);    B(600);    B(1200);   B(1800);
        B(2400);   B(4800);   B(9600);   B(19200);  B(38400);
        B(57600);  B(115200); B(230400); B(460800); B(500000);
        B(576000); B(921600); B(1000000);B(1152000);B(1500000);
default: return 0;
}
#undef B
}

/* Open serial port in raw mode, with custom baudrate if necessary */
static int serial_open(const char *device, int rate)
{
	struct termios options;
	int fd;
	int speed = 0;

	/* Open and configure serial port */
	if ((fd = open(device,O_RDWR|O_NOCTTY)) == -1)
		return -1;

	speed = rate_to_constant(rate);

#ifndef __MACH__
	if (speed == 0) {
		/* Custom divisor */
        struct serial_struct serinfo;
		serinfo.reserved_char[0] = 0;
		if (ioctl(fd, TIOCGSERIAL, &serinfo) < 0)
			return -1;
		serinfo.flags &= ~ASYNC_SPD_MASK;
		serinfo.flags |= ASYNC_SPD_CUST;
		serinfo.custom_divisor = (serinfo.baud_base + (rate / 2)) / rate;
		if (serinfo.custom_divisor < 1)
			serinfo.custom_divisor = 1;
		if (ioctl(fd, TIOCSSERIAL, &serinfo) < 0)
			return -1;
		if (ioctl(fd, TIOCGSERIAL, &serinfo) < 0)
			return -1;
		if (serinfo.custom_divisor * rate != serinfo.baud_base) {
			warnx("actual baudrate is %d / %d = %f",
			      serinfo.baud_base, serinfo.custom_divisor,
			      (float)serinfo.baud_base / serinfo.custom_divisor);
		}
	}
#endif

	fcntl(fd, F_SETFL, 0);
	tcgetattr(fd, &options);
	cfsetispeed(&options, speed ?: 460800);
	cfsetospeed(&options, speed ?: 460800);
	cfmakeraw(&options);
	options.c_cflag |= (CLOCAL | CREAD);
	options.c_cflag &= ~CRTSCTS;
	if (tcsetattr(fd, TCSANOW, &options) != 0)
		return -1;

	return fd;
}

static int tofu(const char *password)
{
    int ret = 0, i;
    int outlen = CRYPTO_KEY_SIZE;
    int siglen = PK_SIGNATURE_SIZE;
    uint8_t clr_sig[PK_SIGNATURE_SIZE];
    struct tofu_packet tofu;
    uint8_t Ke[CRYPTO_KEY_SIZE], Sm[CRYPTO_KEY_SIZE];
    uint8_t resp[256];
    char privkey_fname[32];
    uint32_t mgc = 0xffffffff;
    int hash, hash_len;
    uint8_t iv[24];
    uint8_t sig_dec_verify[64];
    struct pollfd pfd;
    uint32_t qxLen = CRYPTO_KEY_SIZE, qyLen = CRYPTO_KEY_SIZE,
             dLen = CRYPTO_KEY_SIZE;


    tofu.cdc_magic = VAULT_MAGIC;
    tofu.cdc_cmd = CDC_TOFU_INIT;
    tofu.cdc_len = sizeof(struct tofu_packet) - 8;
    
    tofu.vh.magic = VAULT_MAGIC;
    tofu.vh.size = sizeof(struct tofu_packet) - 8;
    

    ecc_key ecc;
    ChaCha cha;
    RNG rng;
    
    Sha512 sha;


    /* Initialize Rng (needed for master key + salt + secret */
    wc_InitRng(&rng);
    
    /* Generate ID */
    wc_RNG_GenerateBlock(&rng, (void *)tofu.vh.id, 8);
    sprintf(privkey_fname, "%08x-%08x.der", tofu.vh.id[0], tofu.vh.id[1]);
    printf("ID: %s\n", privkey_fname);

    /* Generate salt */
    wc_RNG_GenerateBlock(&rng, tofu.vh.host_salt, SALT_LEN);
    printf("Salt: ");
    for (i = 0; i < SALT_LEN; i++) {
        printf("%02x", tofu.vh.host_salt[i]);
    }
    printf("\n");

    /* Derive device key (password based) to create encryption key */
    ret = wc_PBKDF2(Ke, password, strlen(password), tofu.vh.host_salt, SALT_LEN, 1000, CRYPTO_KEY_SIZE, WC_SHA512);
    wc_Chacha_SetKey(&cha, Ke, CRYPTO_KEY_SIZE);
    wc_Chacha_SetIV(&cha, tofu.vh.host_seed, 0);

    /* Generate master signing key */
    wc_ecc_init(&ecc);
    wc_ecc_make_key(&rng, CRYPTO_KEY_SIZE, &ecc);
    qxLen = CRYPTO_KEY_SIZE;
    qyLen = CRYPTO_KEY_SIZE;
    wc_ecc_export_public_raw(&ecc, tofu.vh.auth_pubkey, &qxLen, tofu.vh.auth_pubkey + CRYPTO_KEY_SIZE, &qyLen);
    printf("Public signing key %d: ", qxLen + qyLen);
    for (i = 0; i < qxLen + qyLen; i++) {
        printf("%02x", tofu.vh.auth_pubkey[i]);
    }
    printf("\n");

    
    /* Generate key seed */
    wc_RNG_GenerateBlock(&rng, tofu.vh.host_seed, SEED_LEN);
    printf("Seed: ");
    for (i = 0; i < SEED_LEN; i++) {
        printf("%02x", tofu.vh.host_seed[i]);
    }
    printf("\n");

    wc_InitSha512(&sha);
    for (hash = 0; hash < SHA_PAYLOAD_SIZE;) {
        hash_len = WC_SHA512_BLOCK_SIZE;
        if ((SHA_PAYLOAD_SIZE - hash) < hash_len)
            hash_len = SHA_PAYLOAD_SIZE - hash;
        wc_Sha512Update(&sha, ((uint8_t*)&tofu.vh) + hash, hash_len);
        hash += hash_len;
        printf("update %d/%d\n", hash, SHA_PAYLOAD_SIZE);
    }
    wc_Sha512Final(&sha, tofu.vh.digest);
    printf("Digest (%d): ", sizeof(tofu.vh.digest));
    for (i = 0; i < sizeof(tofu.vh.digest); i++) {
        printf("%02x ", tofu.vh.digest[i]);
        if ((i % 16) == 15)
            printf("\n");
    }
    printf("\n");
    wc_Sha512Free(&sha);

    {
        mp_int r, s;
        mp_init(&r); mp_init(&s);
        ret = wc_ecc_sign_hash_ex(tofu.vh.digest, VAULT_DIGEST_SIZE, &rng, &ecc,
                &r, &s); 
        mp_to_unsigned_bin(&r, &clr_sig[0]);
        mp_to_unsigned_bin(&s, &clr_sig[CRYPTO_KEY_SIZE]);
        mp_clear(&r);
        mp_clear(&s);
    }
    printf("Signature (%d) returned %d:\n", PK_SIGNATURE_SIZE, ret);
    for (i = 0; i < PK_SIGNATURE_SIZE; i++) {
        printf("%02x ", clr_sig[i]);
        if ((i % 16) == 15)
            printf("\n");
    }
    printf("\n");
    {
        int vr, vres;
        mp_int r, s;
        mp_init(&r); mp_init(&s);
        mp_read_unsigned_bin(&r, &clr_sig[0], CRYPTO_KEY_SIZE);
        mp_read_unsigned_bin(&s, &clr_sig[CRYPTO_KEY_SIZE], CRYPTO_KEY_SIZE);
        vr = wc_ecc_verify_hash_ex(&r, &s, tofu.vh.digest, VAULT_DIGEST_SIZE,
            &vres, &ecc);
        mp_clear(&r);
        mp_clear(&s);
        printf("sanity check verify ret %d res %d\n", vr, vres);
    }


    wc_Chacha_SetIV(&cha, tofu.vh.host_seed, 0);
    wc_Chacha_Process(&cha, tofu.vh.signature, clr_sig, siglen);
    wc_Chacha_SetIV(&cha, tofu.vh.host_seed, 0);
    wc_Chacha_Process(&cha, sig_dec_verify, tofu.vh.signature, siglen);

    if (memcmp(clr_sig, sig_dec_verify, siglen) != 0) {
        printf("Error verifying chacha encryption\n");
        return 4;
    }

    printf("Packet (%d): ", sizeof(tofu));
    for (i = 0; i < sizeof(tofu); i++) {
        printf("%02x ", *(((uint8_t*)(&tofu)) + i));
        if ((i % 16) == 15)
            printf("\n");
    }
    printf("\n");

    pfd.fd = sfd;
    pfd.events = POLLIN;
    for(i = 0; i < 10; i++) {
        char buf[8];
        ret = poll(&pfd, 1, 100);
        if (ret > 0) {
            read(sfd, buf, 8);
        } 
    }

    ret = write(sfd, &tofu, sizeof(tofu));
    printf("Sent %d/%d bvtes (tofu) \r\n", ret, sizeof(tofu));
    do {
        ret = read(sfd, &mgc, 4);
    } while (mgc != 0x5afeca5e);

    memcpy (resp, &mgc, 4);
    ret = read(sfd, resp + 4, 64 + 4);
    if (ret >= 4) {
        uint16_t code = *((uint16_t *)(resp + 4));
        uint16_t len = *((uint16_t *)(resp + 6));

        printf("received %d bytes: resp %04x len %04x\n", ret + 4, *((uint16_t *)(resp + 4)), *((uint16_t *)(resp + 6)));
        if (ret > 4) {
            printf("MSG: %s\r\n", resp + 8);
        }
        if (code == CDC_OK) {
            int fd;
            uint8_t privkey_buf[2 * CRYPTO_KEY_SIZE];
            char pkfname[MAX_PATH];
            char *homedir;
            homedir = getenv("HOME");
            if (!homedir || strlen(homedir) == 0) {
                perror("getenv(HOME)");
                exit(1);
            }

            snprintf(pkfname, MAX_PATH - 1, "%s/.pvault/%s", homedir, privkey_fname);
            printf("Saving private key %s\n", privkey_fname);
            fd = open(pkfname, O_WRONLY | O_CREAT | O_EXCL, 0600);
            if (fd < 0) {
                printf("FATAL Error opening private key: %s\n", strerror(errno));
                return 6;
            }
            outlen = 2 * CRYPTO_KEY_SIZE;
            if (wc_ecc_export_private_raw(&ecc, privkey_buf, &qxLen,
                        privkey_buf + qxLen, &qyLen,
                        privkey_buf + qxLen + qyLen, &dLen) != 0) {
                fprintf(stderr, "Unable to export private key to DER\n");
                exit(2);
            }
            outlen = qxLen + qyLen + dLen;
            printf("Exporting private key (%d bytes)\n", outlen);
            write(fd, privkey_buf, outlen); 
            wc_ecc_free(&ecc);
            close(fd);
        }
    }
    return 0;
}

static void addservice(struct vault_service *vs, uint8_t *Ke)
{
    struct vault_service vs_local;
    char gadget_passphrase[64];
    int ret;
    uint8_t Ke_local[CRYPTO_KEY_SIZE];
    uint8_t eccRawKey[CRYPTO_KEY_SIZE * 3];
    uint8_t snd_buffer[sizeof(struct cdc_packet_hdr) + SVC_SIZE];
    struct cdc_packet_hdr *hdr = (struct cdc_packet_hdr *)snd_buffer;
    int i;
    struct pollfd pfd;
    uint32_t mgc = 0;
    int interactive = 0;
    uint32_t slot = 0;
    Sha512 sha;
    ecc_key ecc;
    ChaCha cha;
    RNG rng;
    strcpy(pass_prompt, pass_prompt_svc);
    strcpy(pass_confirm, pass_confirm_svc);


    if (!vs) {
        vs = &vs_local;
        Ke = Ke_local;
        interactive = 1;
        printf("Service name (max 16 char.) : ");
        scanf("%s", vs->name); 
        printf("\n");
        printf("Username (empty = none) : ");
        if (scanf("%s", vs->user) <= 0)
            vs->user[0] = '\0';
        getchar();
        while(!askpass()) {
            printf("Press a key to continue...\r\n");
            getchar();
        }
        strcpy((char *)(vs->pass), passphrase); 
    }
    vs->flags = SVC_FLAG_ACTIVE;
    vs->reserved = 0;

    if (!Ke) {
        printf("Unspecified key\n");
        return;
    }


    wc_InitSha512(&sha);
    wc_Sha512Update(&sha, (uint8_t *)vs, WC_SHA512_BLOCK_SIZE);
    wc_Sha512Final(&sha, vs->dig);

    printf("\n----- RECAP -----\n");
    printf("Service name: %s\n", vs->name);
    printf("Username: %s\n", vs->user);
    printf("Password: ******\n");
    printf("Sha: ");
    for (i = 0; i < VAULT_DIGEST_SIZE; i++) {
        printf("%02x", vs->dig[i]);
    }
    printf("\n-----------------\n\n");


    if (interactive)  {
        printf("Service configured, press a key to send, ctrl+c to abort...\r\n");
        disableEcho();
        printf("Enter gadget passphrase to encrypt and send info: ");
        getPassword(gadget_passphrase, 64);
        enableEcho();
        printf("\n");
        /* Derive device key (password based) to create encryption key */
        ret = wc_PBKDF2(Ke, gadget_passphrase, strlen(gadget_passphrase),
                Status.salt, SALT_LEN, 1000, CRYPTO_KEY_SIZE, WC_SHA512);

        if (ret < 0) {
            printf("Failed to derive password.\n");
            return;
        }
    }


    lseek(pk_fd, SEEK_SET, 0);
    ret = read(pk_fd, eccRawKey, CRYPTO_KEY_SIZE * 3);
    if (ret != CRYPTO_KEY_SIZE * 3) {
        printf("Error reading signing key\n");
        return;
    }

    wc_ecc_init(&ecc);
    if (wc_ecc_import_unsigned(&ecc, eccRawKey, eccRawKey + CRYPTO_KEY_SIZE,
              eccRawKey + 2 * CRYPTO_KEY_SIZE, ECC_SECP256R1) < 0) { 
        printf("Failed importing signing key\n");
        return;
    }

    slot = Status.first_avail;

    /* Initialize Rng (needed for sign)*/
    wc_InitRng(&rng);

    printf("Slot is %u\n", slot);
    {
        mp_int r, s;
        mp_init(&r); mp_init(&s);
        ret = wc_ecc_sign_hash_ex(vs->dig, VAULT_DIGEST_SIZE, &rng, &ecc,
                &r, &s);
        mp_to_unsigned_bin(&r, &vs->sig[0]);
        mp_to_unsigned_bin(&s, &vs->sig[CRYPTO_KEY_SIZE]);
        mp_clear(&r);
        mp_clear(&s);
    }

    printf("Signature (%d):\n", PK_SIGNATURE_SIZE);
    for (i = 0; i < PK_SIGNATURE_SIZE; i++) {
        printf("%02x ", vs->sig[i]);
        if ((i % 16) == 15)
            printf("\n");
    }
    printf("\n");

    wc_ecc_free(&ecc);

    /* Copy the part in clear (flags, reserved) */
    memcpy(snd_buffer + sizeof(struct cdc_packet_hdr), vs, SVC_ENC_OFF);

    /* Set Key and IV */
    wc_Chacha_SetKey(&cha, Ke, CRYPTO_KEY_SIZE);
    wc_Chacha_SetIV(&cha, Status.seed, SVC_FLASH_OFFSET + slot * SVC_SIZE);
    wc_Chacha_Process(&cha, snd_buffer + sizeof(struct cdc_packet_hdr) + SVC_ENC_OFF, 
            ((uint8_t *)(vs)) + SVC_ENC_OFF, SVC_ENC_SIZE);

    hdr->magic = VAULT_MAGIC;
    hdr->cmd = CDC_ADDSERV;
    hdr->len = SVC_SIZE;

    pfd.fd = sfd;
    pfd.events = POLLIN;
    for (i = 0; i < 10; i++) {
        char buf[8];
        ret = poll(&pfd, 1, 10);
        if (ret > 0) {
            read(sfd, buf, 8);
        } 
    }
    
    printf("Packet (%d): ", sizeof(snd_buffer));
    for (i = 0; i < sizeof(snd_buffer); i++) {
        printf("%02x ", snd_buffer[i]);
        if ((i % 16) == 15)
            printf("\n");
    }
    printf("\n");

    printf("Sending svc packet... \n");
    write(sfd, snd_buffer, sizeof(snd_buffer));

    for (i = 0; i < 3; i++) {
        char buf[256];
        ret = poll(&pfd, 1, 5000);
        if (ret > 0) {
            uint16_t code;
            do {
                ret = read(sfd, &mgc, 4);
            } while (mgc != 0x5afeca5e);
            ret = read(sfd, buf, 256);
            code = *((uint16_t *)(buf));
           
            printf("Received response (code %02x, len %d): %s\n",
                    code, ret, (code == CDC_FAIL)?(buf + 4):""); 
            if ((code == CDC_FAIL) || ((code == CDC_OK) && ret == 4))
                break;
        } 
    }

}


char* removeQuotes(char* str) {
    int len = strlen(str);
    int i;

    for (i = 0; i < len; i++)
        if (str[i] != '\"')
            break;
    str = str + i;

    len = strlen(str);

    for (i = len -1; i > 0; i--)
        if (str[i] == '\"')
            str[i] = '\0';
    return str;
}

static void parseCSV(const char* filename) {
    struct vault_service svc;
    FILE* file = fopen(filename, "r");
    uint8_t cKe[CRYPTO_KEY_SIZE];
    int ret;
    char gadget_passphrase[64 + 1];
    char line[1024];
    int header_line_skipped = 0;
    if (file == NULL) {
        printf("Failed to open the file.\n");
        return;
    }

    printf("Enter gadget passphrase to encrypt and send info: ");
    disableEcho();
    getPassword(gadget_passphrase, 64);
    enableEcho();
    printf("\n");
    /* Derive device key (password based) to create encryption key */
    ret = wc_PBKDF2(cKe, gadget_passphrase, strlen(gadget_passphrase),
            Status.salt, SALT_LEN, 1000, CRYPTO_KEY_SIZE, WC_SHA512);

    if (ret < 0) {
        printf("Failed to derive password.\n");
        return;
    }

    while (fgets(line, sizeof(line), file)) {
        int i;

        // Skip header line
        if (!header_line_skipped) {
            header_line_skipped = 1;
            continue;
        }

        // Parse each line of the CSV
        char* token = strtok(line, ",");
        if (token == NULL) {
            continue;
        }

        // Remove quotes from the fields
        token = removeQuotes(token);
        char* url = token;

        // Strip "https:" if present
        if (strncmp(url, "https:", 6) == 0) {
            url += 6;
        }
        
        // Strip "http:" if present
        if (strncmp(url, "http:", 5) == 0) {
            url += 6;
        }

        // Strip slashes ("/") if present
        while (*url == '/') {
            url++;
        }

        // Strip content after the domain name if present
        char* slash = strchr(url, '/');
        if (slash != NULL) {
            *slash = '\0';
        }

        // Parse the domain from right to left
        char* domainStart = NULL;
        char domain[MAX_SVC_NAME_LEN + 1] = "";
        char *domainDot = strrchr(url, '.');

        if (domainDot != NULL) {
            *domainDot=',';
            domainStart = strrchr(url, '.');
            *domainDot='.';
        }
        if (domainStart == NULL) {
            strncpy(domain, url, MAX_SVC_NAME_LEN);
            domain[strlen(url)] = '\0';
        } else {
            int domainLen = strlen(domainStart + 1);
            strncpy(domain, domainStart + 1, domainLen);
        }

        // Store the extracted domain, user, and pass fields
        memset(&svc, 0, sizeof(svc));
        printf("Domain: %s\n", domain);
        strncpy(svc.name,domain, 16); 

        token = strtok(NULL, ",");
        if (token != NULL) {
            token = removeQuotes(token);
            strncpy(svc.user, token, sizeof(svc.user) - 1);
        }

        token = strtok(NULL, ",");
        if (token != NULL) {
            token = removeQuotes(token);
            strncpy(svc.pass, token, sizeof(svc.pass));
        }
        for (i = 0; i < 3; i++) { 
            ret = get_vault_status();
            if (ret == 0)
                break;
        }
        if (ret < 0) {
            printf("Adding service: failed.\n");
            getchar();
            return;
        }
        addservice(&svc, cKe);
    }
    fclose(file);
}


void printMenu() {
    printf("\n");
    printf("0) Device status\n");
    printf("1) TOFU\n");
    printf("2) Add service\n");
    printf("3) Import passwords from CSV\n");
    printf("q) Quit\n\n");
}

char getUserInput() {
    char input;
    printf("Choice: ");
    scanf(" %c", &input);
    getchar(); // Consume the newline character
    return input;
}

void processChoice(char choice) {
    char csv_filename[PATH_MAX];
    switch (choice) {
        case '0':
            // Handle option 0 (Device status)
            clearScreen();
            printf("Device status selected.\n");
            break;
        case '1':
            // Handle option 1 (TOFU)
            clearScreen();
            printf("TOFU selected.\n");
            strcpy(pass_prompt, pass_prompt_phrase);
            strcpy(pass_confirm, pass_confirm_phrase);
            memset(passphrase, 0, sizeof(passphrase));
            while(!askpass()) {
                printf("Press a key to continue...\r\n");
                getchar();
            }

            tofu(passphrase);
            break;
        case '2':
            // Handle option 2 (Add service)
            clearScreen();
            printf("Add service selected.\n");
            strcpy(pass_prompt, pass_prompt_phrase);
            strcpy(pass_confirm, pass_confirm_phrase);
            if (Status.state < 0x02) {
                printf("Cannot add service: device not initialized.\n");
                break;
            }
            if (pk_fd < 0) {
                printf("Cannot add service: private key not present for this device.\n");
                break;
            }
            addservice(NULL, NULL);
            break;
        case '3':
            // Handle option 3 (Import CSV)
            clearScreen();
            printf("CSV file: ");
            scanf("%s", csv_filename);
            getchar();
            parseCSV(csv_filename);
            break;
        case 'q':
        case 'Q':
            // Quit the program
            printf("Quitting...\n");
            exit(0);
        default:
            printf("Invalid choice.\n");
            break;
    }

    printf("Press Enter to continue...\n");
    getchar(); // Wait for Enter key
}


static int get_vault_status(void)
{
    uint32_t rcount = 0;
    struct cdc_packet_hdr hdr;
    int ret;
    int i;
    struct pollfd pfd;
    uint8_t rxbuf[sizeof(struct cdc_packet_hdr) + sizeof(struct vault_status)];
    pfd.fd = sfd;
    pfd.events = POLLIN;

    hdr.magic = VAULT_MAGIC;
    hdr.cmd = CDC_STATUS;
    hdr.len = 0;
    for (i = 0; i < 3; i++) {
        ret = write(sfd, &hdr, sizeof(hdr));
        if (ret <= 0)
            return ret;
    
        ret = poll(&pfd, 1, 1000);
        if (ret == 0) {
            printf("Get status: timeout");
            continue;
        }
        if (ret < 0) {
            perror("get_status(): polling device");
            return -1;
        }
        do {
            ret = read(sfd, rxbuf + rcount, sizeof(rxbuf) - rcount);
            if (ret < 0) {
                return -1;
            }
            rcount += ret;
        } while (rcount < sizeof(rxbuf));
        break;
    }
    if (rcount > 0) {
        if (*(uint32_t *)(rxbuf) != VAULT_MAGIC)
            return -2;
        if (*(uint16_t *)(rxbuf + 4) != CDC_STATUS)
            return -3;
        memcpy(&Status, rxbuf + sizeof(struct cdc_packet_hdr), sizeof(Status));
        return 0;
    }
    return -1;
}


static const char vstatenames[][32] = {
    "OFF",
    "TOFU",
    "BOOTUP",
    "VERIFY_PASSPHRASE",
    "VERIFY_FAILED",
    "MAIN_MENU",
    "SETTINGS_MENU",
    "SERVICE_LIST"
};




int main(int argc, char *argv[]) 
{
    clearScreen();
    system("mkdir -p " HOMEPATH_PREFIX);
    
    if (argc != 2) {
        fprintf(stderr, "Usage: %s TTY\n", argv[0]);
        return 2;
    }

    sfd = serial_open(argv[1], 115200);
    if (sfd < 0) {
        perror("serial_open");
        if (errno == EACCES) {
            printf("\n\nPerhaps you should review your permission for %s, or 'sudo adduser $YOUR_USER dialout'\r\n", argv[1]);
        }
        return 1;
    }
    while (1) {
        char choice;
        if (get_vault_status() < 0) {
            printf("Not connected\n");
        } else {
            char pkfname[MAX_PATH];
            char *homedir;
            homedir = getenv("HOME");
            if (!homedir || strlen(homedir) == 0) {
                perror("getenv(HOME)");
                exit(1);
            }

            snprintf(pkfname, MAX_PATH - 1, "%s/.pvault/%08x-%08x.der", homedir, Status.id[0], Status.id[1]);
            printf("Device %08x-%08x Connected.\n", Status.id[0], Status.id[1]);
            pk_fd = open(pkfname, O_RDONLY);
            if (pk_fd < 0) {
                perror("open");
                printf(pkfname);
                printf("\n");

                printf("Private key not available.\n");
            } else {
                printf("Private key loaded.\n");
            }
            printf("State: %s - Services: %hu. Spot: %hu\n", 
                    vstatenames[Status.state], Status.services_active, Status.first_avail);
        }
        printMenu();
        choice = getUserInput();
        processChoice(choice);
    }

    return 0;
}