// SPDX-License-Identifier: MIT
// Copyright (C) 2022 Benjamin Meisels

// !!!!!!!!!!!!!!!!!
// This isn't the code running so don't try and change it 
// !!!!!!!!!!!!!!!!!

#include "mbedtls/md.h"

#include <cstring>
#include <stdio.h>
#include <Wire.h>
#include <U8g2lib.h>

#define EEPROM_I2C_ADDRESS 0x50
#define SHA256_DIGEST_SIZE 32
static const char G_AUTHENTICATION_MARKER[] = "INTENT";
static const size_t G_AUTHENTICATION_MARKER_SIZE = 7;
static const byte G_EXPECTED_HASH_RESULT[SHA256_DIGEST_SIZE] = { // "INTENT"
  0xc6, 0xa5, 0x99, 0x4d,
  0xde, 0xd4, 0xe6, 0xa7,
  0x7c, 0x8a, 0x95, 0x91,
  0x62, 0x1a, 0xa1, 0xbd,
  0x92, 0x17, 0xa4, 0x3b,
  0x02, 0x9d, 0xbd, 0x7f,
  0xaf, 0x8b, 0x2a, 0x16,
  0xe9, 0xb6, 0x81, 0xc1
};

U8X8_SSD1306_128X64_NONAME_HW_I2C display(U8X8_PIN_NONE);

bool is_device_available(uint8_t address)
{
  Wire.beginTransmission(address);
  uint8_t error = Wire.endTransmission();

  return (error == 0);
}

void wait_for_device(uint8_t i2c_device_address) {
  while(!is_device_available(i2c_device_address));
}

void eeprom_write_byte(uint8_t i2c_device_address, uint16_t eeprom_data_address, uint8_t data) {
  Wire.beginTransmission(i2c_device_address);
  Wire.write((uint8_t)(eeprom_data_address >> 8)); //writes the MSB
  Wire.write((uint8_t)(eeprom_data_address & 0xFF)); //writes the LSB
  Wire.write(data);
  Wire.endTransmission();
  wait_for_device(i2c_device_address);
}

byte eeprom_read_byte(uint8_t i2c_device_address, uint16_t eeprom_data_address) {
  uint8_t rdata = 0xFF;
  
  Wire.beginTransmission(i2c_device_address);
  Wire.write((uint8_t)(eeprom_data_address >> 8)); //writes the MSB
  Wire.write((uint8_t)(eeprom_data_address & 0xFF)); //writes the LSB
  Wire.endTransmission();
  Wire.requestFrom(i2c_device_address,(uint8_t)1);
  if (Wire.available()) {
    rdata = Wire.read();
  }
  return rdata;
}

String eeprom_read_string(uint16_t eeprom_data_address, unsigned int size) {
  String result;
  result.reserve(size);
  for (unsigned int i = 0; i < size; ++i) {
      result.concat((char)eeprom_read_byte(EEPROM_I2C_ADDRESS, i + eeprom_data_address));
  } 

  return result;
}

bool SHA1_verify(unsigned char * data, size_t data_size) {
  byte shaResult[SHA256_DIGEST_SIZE] = {0};

  mbedtls_md_context_t ctx;
  
  mbedtls_md_init(&ctx);
  mbedtls_md_setup(&ctx, mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 0);
  mbedtls_md_starts(&ctx);
  mbedtls_md_update(&ctx, data, data_size);
  mbedtls_md_finish(&ctx, shaResult);
  mbedtls_md_free(&ctx);

  if(0 != std::memcmp(shaResult, G_EXPECTED_HASH_RESULT, SHA256_DIGEST_SIZE)) { 
    Serial.println("- Verification failed");
    Serial.print("Expected Hash: ");
    for(int i= 0; i< SHA256_DIGEST_SIZE; i++)
    {
      printf("%02x", G_EXPECTED_HASH_RESULT[i]);
    }
    printf("\n");

    Serial.print("Calculated Hash: ");
    for(int i= 0; i< SHA256_DIGEST_SIZE; i++)
    {
       printf("%02x", shaResult[i]);
    }
    printf("\n");

    return false;
  }
  Serial.println("- Verification successful");
  return true;
}

void setup()
{
  Wire.begin();

  Serial.begin(9600);
  while (!Serial);

  if (!is_device_available(EEPROM_I2C_ADDRESS))
  {
    printf("No device was found\n");
    while(true);
  }

  display.begin();
  display.setPowerSave(0);
  display.clearDisplay();
  display.setCursor(0,0);
  display.setFont(u8x8_font_pxplusibmcgathin_f);

  String authentication_marker = eeprom_read_string(0x4000, G_AUTHENTICATION_MARKER_SIZE);
  Serial.println("- Reading authentication marker");
  Serial.println(authentication_marker);
  
  byte authentication_marker_buffer[G_AUTHENTICATION_MARKER_SIZE] = {'\0'};
  authentication_marker.getBytes(authentication_marker_buffer, G_AUTHENTICATION_MARKER_SIZE);
  authentication_marker_buffer[G_AUTHENTICATION_MARKER_SIZE - 1] = '\0';
  Serial.println("- Verifying authentication marker");
  bool verified = SHA1_verify(authentication_marker_buffer, G_AUTHENTICATION_MARKER_SIZE);
  
  authentication_marker = eeprom_read_string(0x4000, G_AUTHENTICATION_MARKER_SIZE);
  authentication_marker[G_AUTHENTICATION_MARKER_SIZE - 1] = '\0';
  Serial.println("- Reading authentication marker");
  Serial.println(authentication_marker);
  display.setCursor(0,0);
  display.print(authentication_marker);
  if (verified) {
    display.setCursor(0,3);
    display.print("Genuine!");
    if (authentication_marker != G_AUTHENTICATION_MARKER) {
      Serial.println("- Flag Unlocked");
      // Censored flag.
      // The real flag will be printed when you are successful.
      Serial.println("INTENT{PLACEHOLDER}");
    }
  } else {
    display.setCursor(0,3);
    display.print("Counterfit!");
  }
}

void loop()
{
  delay(1000);
}