#include "DKIMSigner.h"
#include <cctype>
#include <cstring>
#include <ctime>
#include <openssl/bio.h>
#include <openssl/err.h>
#include <openssl/pem.h>
#include <sstream>
#include <stdexcept>

namespace mail {

   DKIMSigner::DKIMSigner(std::string domain, std::string selector, std::string privateKeyPath)
     : domain(domain), selector(selector), privateKey(nullptr) {

      FILE *fp = fopen(privateKeyPath.c_str(), "r");
      if (!fp)
         throw std::runtime_error("DKIMSigner: could not open private key " + privateKeyPath);

      privateKey = PEM_read_PrivateKey(fp, nullptr, nullptr, nullptr);
      fclose(fp);

      if (!privateKey)
         throw std::runtime_error("DKIMSigner: could not parse private key " + privateKeyPath);
   }

   DKIMSigner::~DKIMSigner() {
      if (privateKey)
         EVP_PKEY_free(privateKey);
   }

   // Reduces every run of SP/TAB to a single SP, then trims leading and
   // trailing SP entirely - the relaxed canonicalization rule for
   // whitespace *around* a header field's name/value (RFC 6376 3.4.2),
   // distinct from the body's own version of this rule (canonicalizeBody),
   // which reduces internal runs the same way but does NOT strip a
   // leading run down to nothing, only to one space.
   std::string DKIMSigner::collapseWSP(const std::string &s) {
      std::string result;
      bool lastWasSpace = false;
      for (char c : s) {
         if (c == ' ' || c == '\t') {
            if (!lastWasSpace)
               result += ' ';
            lastWasSpace = true;
         } else {
            result += c;
            lastWasSpace = false;
         }
      }
      size_t start = 0;
      while (start < result.size() && result[start] == ' ')
         ++start;
      size_t end = result.size();
      while (end > start && result[end - 1] == ' ')
         --end;
      return result.substr(start, end - start);
   }

   // RFC 6376 3.4.2 (relaxed header canonicalization), applied to one
   // already-unfolded (name, value) pair: lowercase the name, delete
   // WSP around the colon, collapse+trim WSP within the value, one
   // trailing CRLF.
   std::string DKIMSigner::canonicalizeHeaderField(const std::string &name, const std::string &value) {
      std::string lowerName = name;
      for (char &c : lowerName)
         c = std::tolower((unsigned char)c);
      return lowerName + ":" + collapseWSP(value) + "\r\n";
   }

   // RFC 6376 3.4.4 (relaxed body canonicalization): collapse WSP runs
   // within each line to a single SP (a leading run becomes one SP, not
   // zero - unlike header canonicalization above), strip trailing WSP
   // from each line, then drop trailing empty lines so the body ends in
   // exactly one CRLF (or, if it was empty/all-whitespace to begin
   // with, canonicalizes to a single CRLF - RFC 6376 3.4.3's explicit
   // "null body" case).
   std::string DKIMSigner::canonicalizeBody(const std::string &body) {
      std::vector<std::string> lines;
      size_t pos = 0;
      while (pos < body.size()) {
         size_t next = body.find("\r\n", pos);
         if (next == std::string::npos) {
            lines.push_back(body.substr(pos));
            break;
         }
         lines.push_back(body.substr(pos, next - pos));
         pos = next + 2;
      }

      for (auto &line : lines) {
         std::string collapsed;
         bool lastWasSpace = false;
         for (char c : line) {
            if (c == ' ' || c == '\t') {
               if (!lastWasSpace)
                  collapsed += ' ';
               lastWasSpace = true;
            } else {
               collapsed += c;
               lastWasSpace = false;
            }
         }
         while (!collapsed.empty() && collapsed.back() == ' ')
            collapsed.pop_back();
         line = collapsed;
      }

      while (!lines.empty() && lines.back().empty())
         lines.pop_back();

      std::string result;
      for (auto &line : lines) {
         result += line;
         result += "\r\n";
      }
      if (result.empty())
         result = "\r\n";
      return result;
   }

   std::vector<DKIMSigner::HeaderField> DKIMSigner::parseHeaders(const std::string &message) {
      std::vector<HeaderField> fields;
      size_t headerEnd = message.find("\r\n\r\n");
      std::string headerBlock = (headerEnd == std::string::npos) ? message : message.substr(0, headerEnd);

      std::vector<std::string> rawLines;
      size_t pos = 0;
      while (pos < headerBlock.size()) {
         size_t next = headerBlock.find("\r\n", pos);
         std::string line = (next == std::string::npos) ? headerBlock.substr(pos) : headerBlock.substr(pos, next - pos);
         rawLines.push_back(line);
         if (next == std::string::npos)
            break;
         pos = next + 2;
      }

      // RFC 5322 3.2.2 folding: a continuation line starts with SP/TAB
      // and belongs to the previous header field - unfolded here by
      // joining with a single space, matching what canonicalizeHeaderField's
      // own WSP collapsing expects to receive.
      for (auto &line : rawLines) {
         if (line.empty())
            continue;
         if ((line[0] == ' ' || line[0] == '\t') && !fields.empty()) {
            fields.back().value += " " + line;
            continue;
         }
         size_t colon = line.find(':');
         if (colon == std::string::npos)
            continue;
         HeaderField f;
         f.name = line.substr(0, colon);
         f.value = line.substr(colon + 1);
         fields.push_back(f);
      }

      return fields;
   }

   std::string DKIMSigner::bodyOf(const std::string &message) {
      size_t headerEnd = message.find("\r\n\r\n");
      if (headerEnd == std::string::npos)
         return "";
      return message.substr(headerEnd + 4);
   }

   std::string DKIMSigner::base64(const std::string &raw) {
      // BIO-based rather than reusing coreutils::Base64 here deliberately
      // - this file only needs OpenSSL's own headers (already required
      // for EVP signing), so encoding through the same library avoids
      // pulling CoreUtils' ZString/MString machinery into a class whose
      // only job is talking to libcrypto.
      BIO *b64 = BIO_new(BIO_f_base64());
      BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
      BIO *mem = BIO_new(BIO_s_mem());
      BIO_push(b64, mem);
      BIO_write(b64, raw.data(), (int)raw.size());
      BIO_flush(b64);
      char *encodedData = nullptr;
      long encodedLen = BIO_get_mem_data(mem, &encodedData);
      std::string result(encodedData, encodedLen);
      BIO_free_all(b64);
      return result;
   }

   std::string DKIMSigner::sha256Base64(const std::string &data) {
      unsigned char hash[EVP_MAX_MD_SIZE];
      unsigned int hashLen = 0;
      EVP_Digest(data.data(), data.size(), hash, &hashLen, EVP_sha256(), nullptr);
      return base64(std::string((char *)hash, hashLen));
   }

   std::string DKIMSigner::rsaSha256SignBase64(const std::string &data) {
      EVP_MD_CTX *mdctx = EVP_MD_CTX_new();
      std::string result;
      if (EVP_DigestSignInit(mdctx, nullptr, EVP_sha256(), nullptr, privateKey) == 1 &&
          EVP_DigestSignUpdate(mdctx, data.data(), data.size()) == 1) {
         size_t sigLen = 0;
         EVP_DigestSignFinal(mdctx, nullptr, &sigLen);
         std::vector<unsigned char> sig(sigLen);
         if (EVP_DigestSignFinal(mdctx, sig.data(), &sigLen) == 1)
            result = base64(std::string((char *)sig.data(), sigLen));
      }
      EVP_MD_CTX_free(mdctx);
      if (result.empty())
         throw std::runtime_error("DKIMSigner: RSA-SHA256 signing failed");
      return result;
   }

   std::string DKIMSigner::sign(const std::string &message) {

      std::vector<HeaderField> headers = parseHeaders(message);

      // Sign whichever of these standard headers the message actually
      // has - From is the one DMARC alignment actually depends on, the
      // rest are along for integrity (a receiver seeing a signed
      // Subject/Date/Message-ID that doesn't match the visible one is
      // exactly the tampering DKIM exists to catch). Deliberately not
      // signing headers this relay itself can still legitimately need
      // to touch or add downstream of this point.
      static const char *toSign[] = {"From", "To", "Subject", "Date", "Message-ID"};

      std::string signedHeaderNames;
      std::string headerCanon;
      bool first = true;
      for (const char *wanted : toSign) {
         for (auto &f : headers) {
            bool matches = f.name.size() == strlen(wanted);
            if (matches)
               for (size_t i = 0; i < f.name.size() && matches; ++i)
                  matches = std::tolower((unsigned char)f.name[i]) == std::tolower((unsigned char)wanted[i]);
            if (matches) {
               headerCanon += canonicalizeHeaderField(f.name, f.value);
               if (!first)
                  signedHeaderNames += ":";
               signedHeaderNames += wanted;
               first = false;
               break;
            }
         }
      }

      std::string bodyHash = sha256Base64(canonicalizeBody(bodyOf(message)));

      std::stringstream tagsNoSig;
      tagsNoSig << "v=1; a=rsa-sha256; c=relaxed/relaxed; d=" << domain << "; s=" << selector
                << "; t=" << (long)time(nullptr) << "; h=" << signedHeaderNames
                << "; bh=" << bodyHash << "; b=";

      // The signature covers the selected headers (in the order signed
      // above) followed by this DKIM-Signature field itself, with an
      // empty b= - unlike every other signed header, this last one is
      // included WITHOUT a trailing CRLF (RFC 6376 3.7): the field
      // being created doesn't have one yet, since nothing follows it
      // until sign() returns.
      std::string signingInput = headerCanon + canonicalizeHeaderField("DKIM-Signature", tagsNoSig.str());
      signingInput.erase(signingInput.size() - 2); // drop the trailing CRLF just for this last field

      std::string signature = rsaSha256SignBase64(signingInput);

      return "DKIM-Signature: " + tagsNoSig.str() + signature + "\r\n";
   }

}