#include "SPFChecker.h"
#include <algorithm>
#include <arpa/inet.h>
#include <arpa/nameser.h>
#include <cctype>
#include <cstring>
#include <netdb.h>
#include <resolv.h>
#include <sstream>

namespace mail {

   namespace {

      struct Term {
         char qualifier;         // '+', '-', '~', '?'
         std::string name;       // "all", "ip4", "a", "mx", "include", "exists", "redirect", "exp", or an unknown modifier's key
         std::string value;      // everything after ':'/'=' (domain-spec, ip/cidr, etc.) - empty if none
         bool isModifier;        // "redirect=" / "exp=" / unknown "key=value" vs. an actual mechanism
      };

      std::vector<Term> parseTerms(const std::string &record) {
         std::vector<Term> terms;
         std::istringstream stream(record);
         std::string token;
         bool first = true;
         while (stream >> token) {
            if (first) {
               first = false;
               continue; // "v=spf1" itself
            }
            Term t;
            t.qualifier = '+';
            std::string body = token;
            if (!body.empty() && (body[0] == '+' || body[0] == '-' || body[0] == '~' || body[0] == '?')) {
               t.qualifier = body[0];
               body = body.substr(1);
            }
            size_t eq = body.find('=');
            size_t colon = body.find(':');
            if (eq != std::string::npos && (colon == std::string::npos || eq < colon)) {
               t.isModifier = true;
               t.name = body.substr(0, eq);
               t.value = body.substr(eq + 1);
            } else {
               t.isModifier = false;
               size_t sep = body.find_first_of(":/");
               if (sep == std::string::npos) {
                  t.name = body;
                  t.value = "";
               } else {
                  t.name = body.substr(0, sep);
                  t.value = body.substr(body[sep] == ':' ? sep + 1 : sep); // keep leading '/' for cidr-only forms
               }
            }
            for (char &c : t.name)
               c = std::tolower((unsigned char)c);
            terms.push_back(t);
         }
         return terms;
      }

      // Splits "domain[/v4prefix[/v6prefix]]" (the tail of a/mx after an
      // optional ":domain") into its domain part (possibly empty, meaning
      // "use the checked domain") and the two optional prefix lengths
      // (-1 = "not given", meaning "match the full address").
      void splitDomainAndCidr(const std::string &value, std::string &domain, int &v4prefix, int &v6prefix) {
         domain = value;
         v4prefix = -1;
         v6prefix = -1;
         size_t slash = domain.find('/');
         if (slash == std::string::npos)
            return;
         std::string cidrPart = domain.substr(slash + 1);
         domain = domain.substr(0, slash);
         size_t slash2 = cidrPart.find('/');
         std::string v4 = (slash2 == std::string::npos) ? cidrPart : cidrPart.substr(0, slash2);
         std::string v6 = (slash2 == std::string::npos) ? "" : cidrPart.substr(slash2 + 1);
         if (!v4.empty())
            v4prefix = atoi(v4.c_str());
         if (!v6.empty())
            v6prefix = atoi(v6.c_str());
      }

      SPFChecker::Result qualifierToResult(char q) {
         switch (q) {
            case '-': return SPFChecker::Result::Fail;
            case '~': return SPFChecker::Result::SoftFail;
            case '?': return SPFChecker::Result::Neutral;
            default:  return SPFChecker::Result::Pass;
         }
      }

   }

   std::vector<std::string> SPFChecker::lookupSPFRecords(const std::string &domain) {
      std::vector<std::string> records;
      unsigned char answer[4096];
      int len = res_query(domain.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;
         // RFC 7208 3.3: a TXT record's RDATA is one or more
         // length-prefixed character-strings, concatenated here into
         // one logical string before checking for the "v=spf1 " prefix.
         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=spf1", 0) == 0 && (lower.size() == 6 || lower[6] == ' '))
            records.push_back(full);
      }
      return records;
   }

   std::vector<std::string> SPFChecker::lookupA(const std::string &domain) {
      std::vector<std::string> ips;
      struct addrinfo hints {};
      hints.ai_family = AF_INET;
      hints.ai_socktype = SOCK_STREAM;
      struct addrinfo *result = nullptr;
      if (getaddrinfo(domain.c_str(), nullptr, &hints, &result) != 0 || !result)
         return ips;
      for (struct addrinfo *rp = result; rp != nullptr; rp = rp->ai_next) {
         char buf[INET6_ADDRSTRLEN];
         void *addr = &((struct sockaddr_in *)rp->ai_addr)->sin_addr;
         if (inet_ntop(AF_INET, addr, buf, sizeof(buf)))
            ips.push_back(buf);
      }
      freeaddrinfo(result);
      return ips;
   }

   std::vector<std::string> SPFChecker::lookupMX(const std::string &domain) {
      std::vector<std::string> hosts;
      unsigned char answer[4096];
      int len = res_query(domain.c_str(), C_IN, T_MX, answer, sizeof(answer));
      if (len <= 0)
         return hosts;
      ns_msg handle;
      if (ns_initparse(answer, len, &handle) != 0)
         return hosts;
      int count = ns_msg_count(handle, ns_s_an);
      // RFC 7208 4.6.4: processing more than 10 MX names per mx
      // mechanism is itself a PermError condition - capped here rather
      // than trusting a (possibly hostile) record to stay reasonable.
      for (int i = 0; i < count && i < 10; ++i) {
         ns_rr rr;
         if (ns_parserr(&handle, ns_s_an, i, &rr) != 0 || ns_rr_type(rr) != ns_t_mx)
            continue;
         char host[NS_MAXDNAME];
         if (dn_expand(ns_msg_base(handle), ns_msg_end(handle), ns_rr_rdata(rr) + 2, host, sizeof(host)) >= 0)
            hosts.push_back(host);
      }
      return hosts;
   }

   bool SPFChecker::ipInCIDR(const std::string &ip, const std::string &cidrSpec, int defaultPrefixV4, int defaultPrefixV6) {
      bool isV6 = ip.find(':') != std::string::npos;

      if (isV6) {
         struct in6_addr a{}, n{};
         if (inet_pton(AF_INET6, ip.c_str(), &a) != 1 || inet_pton(AF_INET6, cidrSpec.c_str(), &n) != 1)
            return false;
         int prefix = (defaultPrefixV6 >= 0) ? defaultPrefixV6 : 128;
         if (prefix < 0 || prefix > 128)
            return false;
         int fullBytes = prefix / 8;
         int remBits = prefix % 8;
         if (memcmp(a.s6_addr, n.s6_addr, fullBytes) != 0)
            return false;
         if (remBits > 0) {
            unsigned char mask = (unsigned char)(0xFF << (8 - remBits));
            if ((a.s6_addr[fullBytes] & mask) != (n.s6_addr[fullBytes] & mask))
               return false;
         }
         return true;
      } else {
         struct in_addr a{}, n{};
         if (inet_pton(AF_INET, ip.c_str(), &a) != 1 || inet_pton(AF_INET, cidrSpec.c_str(), &n) != 1)
            return false;
         int prefix = (defaultPrefixV4 >= 0) ? defaultPrefixV4 : 32;
         if (prefix < 0 || prefix > 32)
            return false;
         uint32_t mask = (prefix == 0) ? 0 : htonl(0xFFFFFFFFu << (32 - prefix));
         return (a.s_addr & mask) == (n.s_addr & mask);
      }
   }

   SPFChecker::Result SPFChecker::checkRecursive(const std::string &domain, const std::string &clientIP, int &mechanismCount) {

      std::vector<std::string> records = lookupSPFRecords(domain);
      if (records.empty())
         return Result::None;
      if (records.size() > 1)
         return Result::PermError; // RFC 7208 4.5: more than one v=spf1 record is a PermError

      std::vector<Term> terms = parseTerms(records[0]);
      std::string redirectTarget;

      for (auto &t : terms) {

         if (t.isModifier) {
            if (t.name == "redirect")
               redirectTarget = t.value;
            continue; // "exp" and unrecognized modifiers are not evaluated
         }

         bool matched = false;

         if (t.name == "all") {
            matched = true;

         } else if (t.name == "ip4" || t.name == "ip6") {
            std::string cidr;
            int v4p = -1, v6p = -1;
            splitDomainAndCidr(t.value, cidr, v4p, v6p);
            // ip4/ip6 use a single trailing /prefix, not the dual form -
            // splitDomainAndCidr's "second slash" case does not apply
            // here, so whichever of v4p/v6p got set is the one that
            // matters for this term's own family.
            int prefix = (t.name == "ip4") ? v4p : v6p;
            matched = ipInCIDR(clientIP, cidr, t.name == "ip4" ? prefix : -1, t.name == "ip6" ? prefix : -1);

         } else if (t.name == "a" || t.name == "mx") {
            if (t.value.find('%') != std::string::npos)
               continue; // macro - not supported, skip (see class doc comment)
            std::string targetDomain;
            int v4p = -1, v6p = -1;
            splitDomainAndCidr(t.value, targetDomain, v4p, v6p);
            if (targetDomain.empty())
               targetDomain = domain;

            std::vector<std::string> candidateDomains;
            if (t.name == "a") {
               candidateDomains.push_back(targetDomain);
            } else {
               if (++mechanismCount > 10)
                  return Result::PermError;
               candidateDomains = lookupMX(targetDomain);
            }

            for (auto &cd : candidateDomains) {
               if (++mechanismCount > 10)
                  return Result::PermError;
               for (auto &ip : lookupA(cd)) {
                  if (ipInCIDR(clientIP, ip, v4p, v6p)) {
                     matched = true;
                     break;
                  }
               }
               if (matched)
                  break;
            }

         } else if (t.name == "exists") {
            if (t.value.find('%') != std::string::npos)
               continue; // macro - not supported, skip
            if (++mechanismCount > 10)
               return Result::PermError;
            matched = !lookupA(t.value).empty();

         } else if (t.name == "include") {
            if (++mechanismCount > 10)
               return Result::PermError;
            Result sub = checkRecursive(t.value, clientIP, mechanismCount);
            // RFC 7208 5.2: only a Pass from the included domain makes
            // "include" itself match; Fail/SoftFail/Neutral there just
            // mean "this mechanism didn't match", not that the whole
            // record fails - but an error in the sub-check is a real
            // error here too, not silently swallowed.
            if (sub == Result::PermError || sub == Result::TempError)
               return sub;
            matched = (sub == Result::Pass);

         } else if (t.name == "ptr") {
            continue; // deprecated, unsupported - see class doc comment

         } else {
            // Unrecognized mechanism name - RFC 7208 5: PermError.
            return Result::PermError;
         }

         if (matched)
            return qualifierToResult(t.qualifier);
      }

      if (!redirectTarget.empty()) {
         if (++mechanismCount > 10)
            return Result::PermError;
         return checkRecursive(redirectTarget, clientIP, mechanismCount);
      }

      // RFC 7208 4.7: no mechanism matched and no redirect - Neutral,
      // as if "?all" had been appended.
      return Result::Neutral;
   }

   SPFChecker::Result SPFChecker::check(const std::string &domain, const std::string &clientIP) {
      if (domain.empty())
         return Result::None;
      int mechanismCount = 0;
      return checkRecursive(domain, clientIP, mechanismCount);
   }

   std::string SPFChecker::resultName(Result result) {
      switch (result) {
         case Result::Pass:      return "pass";
         case Result::Fail:      return "fail";
         case Result::SoftFail:  return "softfail";
         case Result::Neutral:   return "neutral";
         case Result::None:      return "none";
         case Result::PermError: return "permerror";
         case Result::TempError: return "temperror";
      }
      return "none";
   }

   std::string SPFChecker::receivedSPFHeader(Result result, const std::string &domain,
                                              const std::string &clientIP, const std::string &helo,
                                              const std::string &receivingHost) {
      std::string name = resultName(result);
      std::string explanation;
      switch (result) {
         case Result::Pass:      explanation = "domain of " + domain + " designates " + clientIP + " as permitted sender"; break;
         case Result::Fail:      explanation = "domain of " + domain + " does not designate " + clientIP + " as permitted sender"; break;
         case Result::SoftFail:  explanation = "domain of " + domain + " does not designate " + clientIP + " as permitted sender (soft fail)"; break;
         case Result::Neutral:   explanation = "domain of " + domain + " has no definitive policy for " + clientIP; break;
         case Result::None:      explanation = "domain of " + domain + " does not designate permitted sender hosts"; break;
         case Result::PermError: explanation = "permanent error in processing domain of " + domain; break;
         case Result::TempError: explanation = "temporary error in processing domain of " + domain; break;
      }
      return "Received-SPF: " + name + " (" + receivingHost + ": " + explanation + ") client-ip=" + clientIP +
             "; envelope-from=" + domain + "; helo=" + helo + ";\r\n";
   }

}