#include "DMARCChecker.h"
#include <arpa/nameser.h>
#include <cctype>
#include <resolv.h>
#include <sstream>

namespace mail {

   namespace {

      std::vector<std::string> lookupDMARCRecords(const std::string &domain) {
         std::vector<std::string> records;
         std::string queryName = "_dmarc." + domain;
         unsigned char answer[4096];
         int len = res_query(queryName.c_str(), C_IN, T_TXT, answer, sizeof(answer));
         if (len <= 0)
            return records;
         ns_msg handle;
         if (ns_initparse(answer, len, &handle) != 0)
            return records;
         int count = ns_msg_count(handle, ns_s_an);
         for (int i = 0; i < count; ++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;
            }
            std::string lower = full;
            for (char &c : lower)
               c = std::tolower((unsigned char)c);
            if (lower.rfind("v=dmarc1", 0) == 0)
               records.push_back(full);
         }
         return records;
      }

      void parseTags(const std::string &record, DMARCChecker::Policy &p) {
         std::istringstream stream(record);
         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();
            for (char &c : key)
               c = std::tolower((unsigned char)c);
            if (key == "p")
               p.policy = value;
            else if (key == "sp")
               p.subdomainPolicy = value;
            else if (key == "adkim")
               p.strictDKIM = (value == "s");
            else if (key == "aspf")
               p.strictSPF = (value == "s");
            else if (key == "pct")
               p.percent = atoi(value.c_str());
         }
         if (p.subdomainPolicy.empty())
            p.subdomainPolicy = p.policy;
      }

   }

   std::string DMARCChecker::organizationalDomain(const std::string &domain) {
      std::vector<std::string> labels;
      std::istringstream stream(domain);
      std::string label;
      while (std::getline(stream, label, '.'))
         if (!label.empty())
            labels.push_back(label);
      if (labels.size() <= 2)
         return domain;
      return labels[labels.size() - 2] + "." + labels[labels.size() - 1];
   }

   bool DMARCChecker::aligned(const std::string &identityDomain, const std::string &fromDomain, bool strict) {
      if (identityDomain.empty() || fromDomain.empty())
         return false;
      auto lower = [](std::string s) { for (char &c : s) c = std::tolower((unsigned char)c); return s; };
      if (strict)
         return lower(identityDomain) == lower(fromDomain);
      return lower(organizationalDomain(identityDomain)) == lower(organizationalDomain(fromDomain));
   }

   bool DMARCChecker::lookupPolicy(const std::string &domain, Policy &policyOut) {
      std::vector<std::string> records = lookupDMARCRecords(domain);
      if (records.size() == 1) {
         parseTags(records[0], policyOut);
         policyOut.found = true;
         return true;
      }
      if (records.size() > 1)
         return false; // RFC 7489 6.6.3: more than one record is treated as none published

      // RFC 7489 6.6.3: no record at the exact domain falls back to the
      // organizational domain's own record (using ITS sp=, or p= if sp=
      // is absent, as the effective policy for this subdomain) - but
      // only when domain actually IS a subdomain of something else, or
      // this would just re-query the same name again.
      std::string org = organizationalDomain(domain);
      if (org == domain)
         return false;
      std::vector<std::string> orgRecords = lookupDMARCRecords(org);
      if (orgRecords.size() != 1)
         return false;
      Policy orgPolicy;
      parseTags(orgRecords[0], orgPolicy);
      policyOut = orgPolicy;
      policyOut.policy = orgPolicy.subdomainPolicy; // sp= (or p=, via parseTags' own fallback) governs here
      policyOut.found = true;
      return true;
   }

   DMARCChecker::Result DMARCChecker::evaluate(const std::string &fromDomain, const std::string &spfDomain,
                                                SPFChecker::Result spfResult,
                                                const std::vector<DKIMVerifier::SignatureResult> &dkimResults,
                                                Policy &appliedPolicy) {
      if (fromDomain.empty() || !lookupPolicy(fromDomain, appliedPolicy))
         return Result::None;

      bool spfAligned = (spfResult == SPFChecker::Result::Pass) && aligned(spfDomain, fromDomain, appliedPolicy.strictSPF);

      bool dkimAligned = false;
      for (auto &d : dkimResults)
         if (d.valid && aligned(d.domain, fromDomain, appliedPolicy.strictDKIM)) {
            dkimAligned = true;
            break;
         }

      return (spfAligned || dkimAligned) ? Result::Pass : Result::Fail;
   }

   std::string DMARCChecker::resultName(Result result) {
      switch (result) {
         case Result::Pass: return "pass";
         case Result::Fail: return "fail";
         case Result::None: return "none";
      }
      return "none";
   }

   std::string DMARCChecker::authenticationResultsHeader(SPFChecker::Result spfResult, const std::string &spfDomain,
                                                          const std::vector<DKIMVerifier::SignatureResult> &dkimResults,
                                                          Result dmarcResult, const std::string &fromDomain,
                                                          const std::string &receivingHost) {
      std::string h = "Authentication-Results: " + receivingHost + ";\r\n";
      h += "\tspf=" + SPFChecker::resultName(spfResult) + " smtp.mailfrom=" + spfDomain + ";\r\n";
      for (auto &d : dkimResults)
         h += "\tdkim=" + std::string(d.valid ? "pass" : "fail") + " header.d=" + d.domain + " header.s=" + d.selector + ";\r\n";
      h += "\tdmarc=" + resultName(dmarcResult) + " header.from=" + fromDomain + ";\r\n";
      return h;
   }

}