#include "ProcessCertificate.h"

#include <iostream>
#include <stdexcept>

#include <openssl/rsa.h>
#include "Utils.h"

using namespace std;

ProcessCertificate::ProcessCertificate(
    const vector<uint8_t>& native_data, const vector<uint8_t>& five_signature,
    const vector<uint8_t>& pa_app_name, long key_id, long flags)
    : native_data_(native_data), pa_app_name_(pa_app_name),
      is_allocated_(false) {
  pa_certificate_ = new PaCertificate_t();
  pa_id_ = Sha256(native_data);
  five_signature_hash_ = Sha256(five_signature);

  PaData& pa_data = pa_certificate_->paData;
  pa_data.fiveSignatureHash.buf = &five_signature_hash_[0];
  pa_data.fiveSignatureHash.size = five_signature_hash_.size();
  pa_data.paId.buf = &pa_id_[0];
  pa_data.paId.size = pa_id_.size();
  pa_data.paVersion = 1;
  pa_data.paFlags = flags;
  pa_data.paKeyId = key_id;
  pa_data.paAppName.buf = &pa_app_name_[0];
  pa_data.paAppName.size = pa_app_name_.size();
}

ProcessCertificate::ProcessCertificate(
    const string& b64_encoded_certificate): pa_certificate_(NULL), is_allocated_(true) {
  PaCertificate *certificate = NULL;
  pa_certificate_ = new PaCertificate_t();

  Asn1DerDecode(Base64Decode(b64_encoded_certificate), &certificate);
  *pa_certificate_ = *certificate;
  free(certificate);
}

ProcessCertificate::ProcessCertificate(
    const string& b64_request,
    const string& b64_signature): is_allocated_(true) {
  PaData *pa_data = NULL;
  Asn1DerDecode(Base64Decode(b64_request), &pa_data);
  pa_certificate_ = new PaCertificate_t();
  pa_certificate_->paData = *pa_data;
  free(pa_data);

  vector<uint8_t> signature = Base64Decode(b64_signature);

  pa_certificate_->signature.buf = (uint8_t *)malloc(signature.size());
  pa_certificate_->signature.size = signature.size();
  memcpy(pa_certificate_->signature.buf, &signature[0], signature.size());
}

string ProcessCertificate::ToAsn1Buffer() {
  return Base64Encode(Asn1DerEncode(pa_certificate_));
}

void ProcessCertificate::Sign(PrivateKey& key) {
  EVP_PKEY_CTX *ctx = NULL;
  vector<uint8_t> md = Sha256(Asn1DerEncode(&pa_certificate_->paData));
  size_t siglen;

  if (!key.GetKey()) {
    throw runtime_error("Sign: EVP_PKEY is null");
  }
  if (!(ctx = EVP_PKEY_CTX_new(key.GetKey(), NULL /* no engine */))) {
    throw runtime_error("Sign: EVP_PKEY_CTX_new failed");
  }
  if (EVP_PKEY_sign_init(ctx) <= 0) {
    throw runtime_error("Sign: EVP_PKEY_sign_init failed");
  }
  if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING) <= 0) {
    throw runtime_error("Sign: EVP_PKEY_CTX_set_rsa_padding failed");
  }
  if (EVP_PKEY_CTX_set_signature_md(ctx, EVP_sha256()) <= 0) {
    throw runtime_error("Sign: EVP_PKEY_CTX_set_signature_md failed");
  }
  if (EVP_PKEY_sign(ctx, NULL, &siglen, &md[0], md.size()) <= 0) {
    throw runtime_error(
        "Sign: EVP_PKEY_sign failed when try to determine buffer length");
  }
  pa_rsa_signature_.resize(siglen);
  if (EVP_PKEY_sign(ctx, &pa_rsa_signature_[0], &siglen, &md[0], md.size()) <= 0) {
    throw runtime_error("Sign: EVP_PKEY_sign failed");
  }

  pa_certificate_->signature.buf = &pa_rsa_signature_[0];
  pa_certificate_->signature.size = pa_rsa_signature_.size();

  if (!Verify(key)) {
    throw runtime_error("Sign: Can not verify signed certificate");
  }
}

bool ProcessCertificate::Verify(KeyInterface& key) const {
  EVP_PKEY_CTX *ctx = NULL;
  vector<uint8_t> md = Sha256(Asn1DerEncode(&pa_certificate_->paData));
  size_t siglen;

  if (!key.GetKey()) {
    throw runtime_error("Verify: EVP_PKEY is null");
  }
  if (!(ctx = EVP_PKEY_CTX_new(key.GetKey(), NULL /* no engine */))) {
    throw runtime_error("Verify: EVP_PKEY_CTX_new failed");
  }
  if (EVP_PKEY_verify_init(ctx) <= 0) {
    throw runtime_error("Verify: EVP_PKEY_verify_init failed");
  }
  if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING) <= 0) {
    throw runtime_error("Verify: EVP_PKEY_CTX_set_rsa_padding failed");
  }
  if (EVP_PKEY_CTX_set_signature_md(ctx, EVP_sha256()) <= 0) {
    throw runtime_error("Verify: EVP_PKEY_CTX_set_signature_md failed");
  }

  int ret = EVP_PKEY_verify(ctx, pa_certificate_->signature.buf,
      pa_certificate_->signature.size, &md[0], md.size());
  if (ret < 0) {
    throw runtime_error("Verify: EVP_PKEY_verify failed");
  }

  return (ret == 1);
}

ProcessCertificate::~ProcessCertificate() {
  if (is_allocated_) {
    Asn1FreeContents(pa_certificate_);
  }
  delete pa_certificate_;
}

string ProcessCertificate::GetRequest() {
  return Base64Encode(Asn1DerEncode(&pa_certificate_->paData));
}

static void DumpOctetString(const string &title, const OCTET_STRING *str) {
  cout << title << ": ";
  cout << hex << uppercase;
  for (size_t x = 0; x < str->size; ++x) {
    cout << (int)str->buf[x];
  }
  cout << dec << "\n";
}

void ProcessCertificate::Dump() const {
  cout << "Certificate version: " << pa_certificate_->paData.paVersion << "\n";
  cout << "App name: " << pa_certificate_->paData.paAppName.buf << "\n";
  DumpOctetString("PA ID", &pa_certificate_->paData.paId);
  DumpOctetString("FIVE signature (sha256)", &pa_certificate_->paData.fiveSignatureHash);
  if (pa_certificate_->paData.paKeyId == 0) {
    cout << "Signed by ENG key\n";
  } else if (pa_certificate_->paData.paKeyId == 1) {
    cout << "Signed by USER key\n";
  } else {
    cout << "Signed by UNKNOWN key\n";
  }
}