keypitecc/main/zk_auth.h

#ifndef ZK_AUTH_H
#define ZK_AUTH_H

#include "cryptoauthlib.h"
#include "host/atca_host.h"
#include "libssh2_config.h"
#include "mbedtls/sha256.h"
#include <cJSON.h>
#include <esp_mac.h>
#include <esp_system.h>
#include <libssh2.h>
#include <mbedtls/aes.h>
#include <mbedtls/base64.h>
#include <mbedtls/ctr_drbg.h>
#include <mbedtls/ecdh.h>
#include <mbedtls/ecp.h>
#include <mbedtls/entropy.h>
#include <mbedtls/md.h>
#include <mbedtls/pkcs5.h>
#include <mbedtls/sha256.h>
#include <stdio.h>
#include <string.h>

// Added network headers for ESP-IDF (lwIP) sockets
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

class ZKAuth {
private:
  mbedtls_ecp_group grp;
  mbedtls_mpi device_private_d;
  mbedtls_ecp_point device_public_Q;
  mbedtls_entropy_context entropy;
  mbedtls_ctr_drbg_context ctr_drbg;

  uint8_t device_public_key[65];    // Uncompressed: 0x04 + X(32) + Y(32)
  uint8_t stored_password_hash[32]; // PBKDF2 hash of the correct password

  bool initialized;

  // Convert binary to hex string
  void bin_to_hex(const uint8_t *bin, size_t bin_len, char *hex) {
    for (size_t i = 0; i < bin_len; i++) {
      sprintf(hex + (i * 2), "%02x", bin[i]);
    }
    hex[bin_len * 2] = '\0';
  }

  // Convert hex string to binary
  bool hex_to_bin(const char *hex, uint8_t *bin, size_t bin_len) {
    if (strlen(hex) != bin_len * 2)
      return false;

    for (size_t i = 0; i < bin_len; i++) {
      if (sscanf(hex + (i * 2), "%2hhx", &bin[i]) != 1) {
        return false;
      }
    }
    return true;
  }

public:
  ZKAuth() : initialized(false) {
    mbedtls_ecp_group_init(&grp);
    mbedtls_mpi_init(&device_private_d);
    mbedtls_ecp_point_init(&device_public_Q);
    mbedtls_entropy_init(&entropy);
    mbedtls_ctr_drbg_init(&ctr_drbg);
  }

  ~ZKAuth() {
    mbedtls_ecp_group_free(&grp);
    mbedtls_mpi_free(&device_private_d);
    mbedtls_ecp_point_free(&device_public_Q);
    mbedtls_entropy_free(&entropy);
    mbedtls_ctr_drbg_free(&ctr_drbg);
  }

  // Initialize the ZK authentication system
  bool init() {
    if (initialized) {
      return true;
    } else {
      initialized = true;
      printf("\n=== ZK Authentication Initialized ===\n");
    }
    return true;
  }

  // Marked as static so it can be passed as a C callback to libssh2
  static int sign_callback(LIBSSH2_SESSION *session, unsigned char **sig,
                           size_t *sig_len, const unsigned char *data,
                           size_t data_len, void **abstract) {
    ATCA_STATUS status;
    uint8_t signature[64];

    // sign the data hash with ATECC608B
    status = atcab_sign(0x0, data, signature);
    if (status != ATCA_SUCCESS) {
      return -1;
    }

    // TODO: implement DER encoding and allocate output blob
    // ...

    return 0;
  }

  // Helper function to write a 32-bit integer in big-endian format
  void write_uint32_be(uint8_t *buf, uint32_t val) {
    buf[0] = (val >> 24) & 0xFF;
    buf[1] = (val >> 16) & 0xFF;
    buf[2] = (val >> 8) & 0xFF;
    buf[3] = val & 0xFF;
  }

  // Function to generate the authorized_keys string
  void generate_ssh_authorized_key(const uint8_t *atecc_pubkey,
                                   uint8_t *out_blob) {
    uint8_t ssh_blob[104];
    uint32_t offset = 0;

    // 1. Key Type
    const char *key_type = "ecdsa-sha2-nistp256";
    write_uint32_be(&ssh_blob[offset], 19);
    offset += 4;
    memcpy(&ssh_blob[offset], key_type, 19);
    offset += 19;

    // 2. Curve Name
    const char *curve_name = "nistp256";
    write_uint32_be(&ssh_blob[offset], 8);
    offset += 4;
    memcpy(&ssh_blob[offset], curve_name, 8);
    offset += 8;

    // 3. Public Key (Uncompressed format: 0x04 + 64 bytes X/Y)
    write_uint32_be(&ssh_blob[offset], 65);
    offset += 4;
    ssh_blob[offset++] = 0x04;
    memcpy(&ssh_blob[offset], atecc_pubkey, 64);
    offset += 64;

    // 4. Base64 Encode the blob
    size_t b64_len = 0;
    // Call once to get required length
    mbedtls_base64_encode(NULL, 0, &b64_len, ssh_blob, 104);

    unsigned char b64_out[b64_len];
    // Call again to actually encode
    mbedtls_base64_encode(b64_out, b64_len, &b64_len, ssh_blob, 104);

    // 5. Print out the final authorized_keys line
    printf("ecdsa-sha2-nistp256 %s esp32-atecc608b\n", b64_out);

    memcpy(out_blob, ssh_blob, 104);
  }

  // Get device identity (for /api/identity endpoint)
  char *get_identity_json() {
    char pubkey_hex[131];     // 65 bytes * 2 + null
    uint8_t atecc_pubkey[64]; // 65 bytes * 2 + null
    uint8_t standard_pubkey[65];
    standard_pubkey[0] = 0x04;
    // Get public key from ATECC608B and convert to hex
    ATCA_STATUS status = atcab_get_pubkey(0, atecc_pubkey);
    if (status != ATCA_SUCCESS) {
      printf("Failed to read public key from ATECC608B: 0x%02X\n", status);
    }

    uint8_t pubkey_blob[256];

    generate_ssh_authorized_key(atecc_pubkey, pubkey_blob);

    // Print the authorized_keys line for debugging
    printf("Generated authorized_keys line:\n");
    for (int i = 0; i < 104; i++) {
      printf("%02x", pubkey_blob[i]);
    }
    printf("\n");

    memcpy(&standard_pubkey[1], atecc_pubkey, 64);
    bin_to_hex(standard_pubkey, 65, pubkey_hex);

    // Get MAC address to use as salt
    uint8_t mac[6];
    esp_read_mac(mac, ESP_MAC_WIFI_STA);
    char mac_hex[13]; // 6 bytes * 2 + null
    bin_to_hex(mac, 6, mac_hex);

    cJSON *root = cJSON_CreateObject();
    cJSON_AddStringToObject(root, "pubKey", pubkey_hex);
    cJSON_AddStringToObject(root, "macAddress", mac_hex);

    char *json_str = cJSON_PrintUnformatted(root);
    cJSON_Delete(root);

    int sock;
    struct sockaddr_in sin;
    const char *hostname = "192.168.4.1";
    unsigned short port = 22;
    const char *username = "key";

    void *my_abstract = NULL;
    LIBSSH2_SESSION *session;

    // Initialize libssh2
    libssh2_init(0);

    // Set up and connect socket
    sock = socket(AF_INET, SOCK_STREAM, 0);
    sin.sin_family = AF_INET;
    sin.sin_port = htons(port);
    inet_pton(AF_INET, hostname, &sin.sin_addr);
    connect(sock, (struct sockaddr *)(&sin), sizeof(struct sockaddr_in));

    // Create SSH session
    session = libssh2_session_init();
    libssh2_session_handshake(session, sock);

    // PUBLIC KEY AUTH: using custom callback
    int rc = libssh2_userauth_publickey(session, username, pubkey_blob,
                                        sizeof(pubkey_blob), sign_callback,
                                        &my_abstract);

    if (rc == 0) {
      printf("SSH authentication successful!\n");
      // ... Continue SSH communication ...
    } else {
      printf("Authentication failed\n");
    }

    // Cleanup
    libssh2_session_disconnect(session, "Bye");
    libssh2_session_free(session);
    close(sock);
    libssh2_exit();

    return json_str;
  }
};

#endif // ZK_AUTH_H