#include "DKIMVerifier.h"
#include <arpa/nameser.h>
#include <cctype>
#include <cstring>
#include <openssl/bio.h>
#include <openssl/evp.h>
#include <openssl/x509.h>
#include <resolv.h>
#include <sstream>

namespace mail {

   namespace {

      struct RawHeader {
         std::string name;         // original case
         std::string rawText;      // exact bytes as they appeared, folding included, ending in CRLF
         std::string unfoldedValue; // continuation lines joined with a single space, not yet WSP-collapsed
      };

      std::vector<RawHeader> parseRawHeaders(const std::string &message) {
         std::vector<RawHeader> headers;
         size_t headerEnd = message.find("\r\n\r\n");
         std::string block = (headerEnd == std::string::npos) ? message : message.substr(0, headerEnd);

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

         for (auto &line : rawLines) {
            if (line.empty())
               continue;
            if ((line[0] == ' ' || line[0] == '\t') && !headers.empty()) {
               headers.back().rawText += line + "\r\n";
               headers.back().unfoldedValue += " " + line;
               continue;
            }
            size_t colon = line.find(':');
            if (colon == std::string::npos)
               continue;
            RawHeader h;
            h.name = line.substr(0, colon);
            h.rawText = line + "\r\n";
            h.unfoldedValue = line.substr(colon + 1);
            headers.push_back(h);
         }
         return headers;
      }

      std::string 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 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) - as in DKIMSigner. Returns rawText
      // unmodified for 3.4.1 (simple), which requires the header field
      // exactly as it appeared, folding included.
      std::string canonicalizeHeaderField(bool simple, const RawHeader &h) {
         if (simple)
            return h.rawText;
         std::string lowerName = h.name;
         for (char &c : lowerName)
            c = std::tolower((unsigned char)c);
         return lowerName + ":" + collapseWSP(h.unfoldedValue) + "\r\n";
      }

      // For synthesizing the DKIM-Signature field itself (with b=
      // emptied) under simple canonicalization, which needs the same
      // "exactly as it appeared" rule applied to a header this code
      // built rather than one that was actually received on the wire -
      // so "exactly as it appeared" here means exactly the tag text,
      // with the SAME folding (none) it already has.
      std::string canonicalizeSyntheticHeader(bool simple, const std::string &name, const std::string &value) {
         if (simple)
            return name + ":" + value + "\r\n";
         std::string lowerName = name;
         for (char &c : lowerName)
            c = std::tolower((unsigned char)c);
         return lowerName + ":" + collapseWSP(value) + "\r\n";
      }

      std::string canonicalizeBody(bool simple, 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;
         }

         if (!simple) {
            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::string base64Decode(const std::string &encoded) {
         BIO *b64 = BIO_new(BIO_f_base64());
         BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
         BIO *mem = BIO_new_mem_buf(encoded.data(), (int)encoded.size());
         BIO_push(b64, mem);
         std::string out(encoded.size(), '\0');
         int len = BIO_read(b64, out.data(), (int)out.size());
         BIO_free_all(b64);
         if (len <= 0)
            return "";
         out.resize(len);
         return out;
      }

      std::string sha1or256Base64(bool useSha1, const std::string &data) {
         unsigned char hash[EVP_MAX_MD_SIZE];
         unsigned int hashLen = 0;
         EVP_Digest(data.data(), data.size(), hash, &hashLen, useSha1 ? EVP_sha1() : EVP_sha256(), nullptr);
         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, hash, (int)hashLen);
         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;
      }

      // Fetches the signer's public key: TXT record at
      // <selector>._domainkey.<domain> (RFC 6376 3.6.1/3.6.2.2), tag
      // format "v=DKIM1; k=rsa; p=<base64 SubjectPublicKeyInfo>" - v=
      // and k= are optional (default to exactly this), only p= is
      // actually required. An empty p= is a published, deliberate key
      // revocation (6376 3.6.1) and correctly fails to load as a key.
      EVP_PKEY *fetchPublicKey(const std::string &selector, const std::string &domain) {
         std::string queryName = selector + "._domainkey." + domain;
         unsigned char answer[4096];
         int len = res_query(queryName.c_str(), C_IN, T_TXT, answer, sizeof(answer));
         if (len <= 0)
            return nullptr;

         ns_msg handle;
         if (ns_initparse(answer, len, &handle) != 0)
            return nullptr;

         std::string record;
         int count = ns_msg_count(handle, ns_s_an);
         for (int i = 0; i < count && record.empty(); ++i) {
            ns_rr rr;
            if (ns_parserr(&handle, ns_s_an, i, &rr) != 0 || ns_rr_type(rr) != ns_t_txt)
               continue;
            const unsigned char *rdata = ns_rr_rdata(rr);
            int rdlen = ns_rr_rdlen(rr);
            std::string full;
            int pos = 0;
            while (pos < rdlen) {
               int segLen = rdata[pos];
               ++pos;
               if (pos + segLen > rdlen)
                  break;
               full.append((const char *)(rdata + pos), segLen);
               pos += segLen;
            }
            if (full.find("p=") != std::string::npos)
               record = full;
         }
         if (record.empty())
            return nullptr;

         std::string p;
         std::istringstream stream(record);
         std::string tag;
         while (std::getline(stream, tag, ';')) {
            size_t start = tag.find_first_not_of(" \t");
            if (start == std::string::npos)
               continue;
            tag = tag.substr(start);
            if (tag.rfind("p=", 0) == 0) {
               p = tag.substr(2);
               // Public key TXT records are often split across several
               // quoted strings by the zone editor for readability -
               // res_query()/ns_parserr() already concatenated the RR's
               // own multiple character-strings, but stray whitespace
               // between them (or trailing junk after the base64) still
               // needs stripping before base64 decoding.
               std::string cleaned;
               for (char c : p)
                  if (!std::isspace((unsigned char)c))
                     cleaned += c;
               p = cleaned;
            }
         }
         if (p.empty())
            return nullptr; // revoked key (6376 3.6.1) or malformed record

         std::string der = base64Decode(p);
         if (der.empty())
            return nullptr;
         const unsigned char *derPtr = (const unsigned char *)der.data();
         return d2i_PUBKEY(nullptr, &derPtr, (long)der.size());
      }

   }

   std::vector<DKIMVerifier::SignatureResult> DKIMVerifier::verify(const std::string &message) {

      std::vector<SignatureResult> results;
      std::vector<RawHeader> headers = parseRawHeaders(message);
      std::string body = bodyOf(message);

      for (size_t sigIndex = 0; sigIndex < headers.size(); ++sigIndex) {
         if (headers[sigIndex].name.size() != 14) continue;
         bool isDkimSig = true;
         {
            static const char *want = "DKIM-Signature";
            for (size_t i = 0; i < 14 && isDkimSig; ++i)
               isDkimSig = std::tolower((unsigned char)headers[sigIndex].name[i]) == std::tolower((unsigned char)want[i]);
         }
         if (!isDkimSig)
            continue;

         SignatureResult result;

         // Parse this DKIM-Signature's own tags.
         std::string tagStr = headers[sigIndex].unfoldedValue;
         std::string a = "rsa-sha256", c = "simple/simple", d, s, h, bh, b;
         std::istringstream stream(tagStr);
         std::string part;
         while (std::getline(stream, part, ';')) {
            size_t start = part.find_first_not_of(" \t");
            if (start == std::string::npos)
               continue;
            part = part.substr(start);
            size_t eq = part.find('=');
            if (eq == std::string::npos)
               continue;
            std::string key = part.substr(0, eq);
            std::string value = part.substr(eq + 1);
            while (!value.empty() && std::isspace((unsigned char)value.back()))
               value.pop_back();
            if (key == "a") a = value;
            else if (key == "c") c = value;
            else if (key == "d") d = value;
            else if (key == "s") s = value;
            else if (key == "h") h = value;
            else if (key == "bh") bh = value;
            else if (key == "b") { std::string cleaned; for (char ch : value) if (!std::isspace((unsigned char)ch)) cleaned += ch; b = cleaned; }
         }

         result.domain = d;
         result.selector = s;

         if (d.empty() || s.empty() || bh.empty() || b.empty()) {
            results.push_back(result); // valid=false: malformed signature
            continue;
         }

         bool simpleHeader = (c.rfind("simple", 0) == 0);
         bool simpleBody = (c.find('/') != std::string::npos) ? (c.substr(c.find('/') + 1).rfind("simple", 0) == 0)
                                                                : simpleHeader; // "c=simple" alone means simple/simple

         bool useSha1 = (a == "rsa-sha1");

         std::string computedBH = sha1or256Base64(useSha1, canonicalizeBody(simpleBody, body));
         if (computedBH != bh) {
            results.push_back(result); // body was altered in transit - valid=false
            continue;
         }

         // Build the signing input the same way DKIMSigner does: every
         // header named in h=, in that order, canonicalized per c='s
         // header algorithm, followed by this DKIM-Signature field
         // itself (with b= emptied) - also canonicalized, and, for
         // simple canonicalization specifically, using the ORIGINAL
         // header exactly as received up through "b=" (nothing after
         // it), not a re-synthesized one, since simple canonicalization
         // forbids changing header text at all.
         std::string signingInput;
         std::istringstream hstream(h);
         std::string wanted;
         while (std::getline(hstream, wanted, ':')) {
            for (auto &hdr : headers) {
               if (hdr.name.size() == wanted.size()) {
                  bool match = true;
                  for (size_t i = 0; i < wanted.size() && match; ++i)
                     match = std::tolower((unsigned char)hdr.name[i]) == std::tolower((unsigned char)wanted[i]);
                  if (match) {
                     signingInput += canonicalizeHeaderField(simpleHeader, hdr);
                     break;
                  }
               }
            }
         }

         size_t bPos = tagStr.rfind("b=");
         std::string tagsUpToB = tagStr.substr(0, bPos + 2);
         std::string dkimFieldCanon;
         if (simpleHeader) {
            // rawText already reads "DKIM-Signature: v=1; ...; b=..." in
            // full (name, colon, and value all together, folding
            // preserved) - simple canonicalization requires using that
            // exact original text unmodified, just truncated right after
            // the literal "b=" so the actual signature value isn't
            // included (RFC 6376 3.7).
            size_t rawBPos = headers[sigIndex].rawText.find("b=");
            dkimFieldCanon = (rawBPos == std::string::npos)
              ? canonicalizeSyntheticHeader(true, headers[sigIndex].name, tagsUpToB)
              : headers[sigIndex].rawText.substr(0, rawBPos + 2);
         } else {
            dkimFieldCanon = canonicalizeSyntheticHeader(false, headers[sigIndex].name, tagsUpToB);
         }
         // The CRLF that would terminate this header field MUST NOT be
         // included (RFC 6376 3.7) - it doesn't exist yet from the
         // signer's perspective. canonicalizeSyntheticHeader's relaxed
         // path always appends one; the simple path (a rawText prefix
         // ending right at "b=") never had one to begin with.
         if (dkimFieldCanon.size() >= 2 && dkimFieldCanon.substr(dkimFieldCanon.size() - 2) == "\r\n")
            dkimFieldCanon.erase(dkimFieldCanon.size() - 2);

         signingInput += dkimFieldCanon;

         EVP_PKEY *pubkey = fetchPublicKey(s, d);
         if (!pubkey) {
            results.push_back(result); // no key published, or revoked - valid=false
            continue;
         }

         std::string sigBytes = base64Decode(b);
         EVP_MD_CTX *mdctx = EVP_MD_CTX_new();
         int rc = 0;
         if (EVP_DigestVerifyInit(mdctx, nullptr, useSha1 ? EVP_sha1() : EVP_sha256(), nullptr, pubkey) == 1)
            rc = EVP_DigestVerify(mdctx, (const unsigned char *)sigBytes.data(), sigBytes.size(),
                                   (const unsigned char *)signingInput.data(), signingInput.size());
         EVP_MD_CTX_free(mdctx);
         EVP_PKEY_free(pubkey);

         result.valid = (rc == 1);
         results.push_back(result);
      }

      return results;
   }

}