/*
 * This file is part of PowerDNS or dnsdist.
 * Copyright -- PowerDNS.COM B.V. and its contributors
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of version 2 of the GNU General Public License as
 * published by the Free Software Foundation.
 *
 * In addition, for the avoidance of any doubt, permission is granted to
 * link this program with OpenSSL and to (re)distribute the binaries
 * produced as the result of such linking.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */
#pragma once

#ifndef DISABLE_DYNBLOCKS
#include <unordered_set>

#include "dolog.hh"
#include "dnsdist-rings.hh"
#include "statnode.hh"

extern "C"
{
#include "dnsdist-lua-inspection-ffi.h"
}

// dnsdist_ffi_stat_node_t is a lightuserdata
template <>
struct LuaContext::Pusher<dnsdist_ffi_stat_node_t*>
{
  static const int minSize = 1;
  static const int maxSize = 1;

  static PushedObject push(lua_State* state, dnsdist_ffi_stat_node_t* ptr) noexcept
  {
    lua_pushlightuserdata(state, ptr);
    return PushedObject{state, 1};
  }
};

using dnsdist_ffi_stat_node_visitor_t = std::function<bool(dnsdist_ffi_stat_node_t*)>;

struct SMTBlockParameters
{
  std::optional<std::string> d_reason;
  std::optional<DNSAction::Action> d_action;
};

struct dnsdist_ffi_stat_node_t
{
  dnsdist_ffi_stat_node_t(const StatNode& node_, const StatNode::Stat& self_, const StatNode::Stat& children_, SMTBlockParameters& blockParameters) :
    node(node_), self(self_), children(children_), d_blockParameters(blockParameters)
  {
  }

  const StatNode& node;
  const StatNode::Stat& self;
  const StatNode::Stat& children;
  SMTBlockParameters& d_blockParameters;
};

using dnsdist_ffi_dynamic_block_inserted_hook = std::function<void(uint8_t type, const char* key, const char* reason, uint8_t action, uint64_t duration, bool warning)>;

class DynBlockRulesGroup
{
private:
  struct Counts
  {
    std::map<uint8_t, uint64_t> d_rcodeCounts;
    std::map<uint16_t, uint64_t> d_qtypeCounts;
    uint64_t queries{0};
    uint64_t responses{0};
    uint64_t respBytes{0};
    uint64_t cacheMisses{0};
  };

  struct DynBlockRule
  {
    DynBlockRule() = default;
    DynBlockRule(const std::string& blockReason, unsigned int blockDuration, unsigned int rate, unsigned int warningRate, unsigned int seconds, DNSAction::Action action) :
      d_blockReason(blockReason), d_blockDuration(blockDuration), d_rate(rate), d_warningRate(warningRate), d_seconds(seconds), d_action(action), d_enabled(true)
    {
    }

    bool matches(const struct timespec& when);
    bool rateExceeded(unsigned int count, const struct timespec& now) const;
    bool warningRateExceeded(unsigned int count, const struct timespec& now) const;

    bool isEnabled() const
    {
      return d_enabled;
    }

    std::string toString() const;

    std::string d_blockReason;
    struct timespec d_cutOff;
    struct timespec d_minTime;
    unsigned int d_blockDuration{0};
    unsigned int d_rate{0};
    unsigned int d_warningRate{0};
    unsigned int d_seconds{0};
    DNSAction::Action d_action{DNSAction::Action::None};
    bool d_enabled{false};
  };

  struct DynBlockRatioRule : DynBlockRule
  {
    DynBlockRatioRule() = default;
    DynBlockRatioRule(const std::string& blockReason, unsigned int blockDuration, double ratio, double warningRatio, unsigned int seconds, DNSAction::Action action, size_t minimumNumberOfResponses) :
      DynBlockRule(blockReason, blockDuration, 0, 0, seconds, action), d_minimumNumberOfResponses(minimumNumberOfResponses), d_ratio(ratio), d_warningRatio(warningRatio)
    {
    }

    bool ratioExceeded(unsigned int total, unsigned int count) const;
    bool warningRatioExceeded(unsigned int total, unsigned int count) const;
    std::string toString() const;

    size_t d_minimumNumberOfResponses{0};
    double d_ratio{0.0};
    double d_warningRatio{0.0};
  };

  struct DynBlockCacheMissRatioRule : public DynBlockRatioRule
  {
    DynBlockCacheMissRatioRule() = default;
    DynBlockCacheMissRatioRule(const std::string& blockReason, unsigned int blockDuration, double ratio, double warningRatio, unsigned int seconds, DNSAction::Action action, size_t minimumNumberOfResponses, double minimumGlobalCacheHitRatio) :
      DynBlockRatioRule(blockReason, blockDuration, ratio, warningRatio, seconds, action, minimumNumberOfResponses), d_minimumGlobalCacheHitRatio(minimumGlobalCacheHitRatio)
    {
    }

    bool checkGlobalCacheHitRatio() const;
    bool ratioExceeded(unsigned int total, unsigned int count) const;
    bool warningRatioExceeded(unsigned int total, unsigned int count) const;
    std::string toString() const;

    double d_minimumGlobalCacheHitRatio{0.0};
  };

  using counts_t = std::unordered_map<AddressAndPortRange, Counts, AddressAndPortRange::hash>;

public:
  DynBlockRulesGroup()
  {
  }

  void setQueryRate(unsigned int rate, unsigned int warningRate, unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action)
  {
    d_queryRateRule = DynBlockRule(reason, blockDuration, rate, warningRate, seconds, action);
  }

  /* rate is in bytes per second */
  void setResponseByteRate(unsigned int rate, unsigned int warningRate, unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action)
  {
    d_respRateRule = DynBlockRule(reason, blockDuration, rate, warningRate, seconds, action);
  }

  void setRCodeRate(uint8_t rcode, unsigned int rate, unsigned int warningRate, unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action)
  {
    auto& entry = d_rcodeRules[rcode];
    entry = DynBlockRule(reason, blockDuration, rate, warningRate, seconds, action);
  }

  void setRCodeRatio(uint8_t rcode, double ratio, double warningRatio, unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action, size_t minimumNumberOfResponses)
  {
    auto& entry = d_rcodeRatioRules[rcode];
    entry = DynBlockRatioRule(reason, blockDuration, ratio, warningRatio, seconds, action, minimumNumberOfResponses);
  }

  void setQTypeRate(uint16_t qtype, unsigned int rate, unsigned int warningRate, unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action)
  {
    auto& entry = d_qtypeRules[qtype];
    entry = DynBlockRule(reason, blockDuration, rate, warningRate, seconds, action);
  }

  void setCacheMissRatio(double ratio, double warningRatio, unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action, size_t minimumNumberOfResponses, double minimumGlobalCacheHitRatio)
  {
    d_respCacheMissRatioRule = DynBlockCacheMissRatioRule(reason, blockDuration, ratio, warningRatio, seconds, action, minimumNumberOfResponses, minimumGlobalCacheHitRatio);
  }

  using smtVisitor_t = std::function<std::tuple<bool, boost::optional<std::string>, boost::optional<int>>(const StatNode&, const StatNode::Stat&, const StatNode::Stat&)>;

  void setSuffixMatchRule(unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action, smtVisitor_t visitor)
  {
    d_suffixMatchRule = DynBlockRule(reason, blockDuration, 0, 0, seconds, action);
    d_smtVisitor = std::move(visitor);
  }

  void setSuffixMatchRuleFFI(unsigned int seconds, const std::string& reason, unsigned int blockDuration, DNSAction::Action action, dnsdist_ffi_stat_node_visitor_t visitor)
  {
    d_suffixMatchRule = DynBlockRule(reason, blockDuration, 0, 0, seconds, action);
    d_smtVisitorFFI = std::move(visitor);
  }

  void setNewBlockHook(const dnsdist_ffi_dynamic_block_inserted_hook& callback)
  {
    d_newBlockHook = callback;
  }

  void setMasks(uint8_t v4, uint8_t v6, uint8_t port)
  {
    d_v4Mask = v4;
    d_v6Mask = v6;
    d_portMask = port;
  }

  void apply()
  {
    struct timespec now;
    gettime(&now);

    apply(now);
  }

  void apply(const struct timespec& now);

  void excludeRange(const Netmask& range)
  {
    d_excludedSubnets.addMask(range);
  }

  void excludeRange(const NetmaskGroup& group)
  {
    d_excludedSubnets.addMasks(group, true);
  }

  void includeRange(const Netmask& range)
  {
    d_excludedSubnets.addMask(range, false);
  }

  void includeRange(const NetmaskGroup& group)
  {
    d_excludedSubnets.addMasks(group, false);
  }

  void removeRange(const Netmask& range)
  {
    d_excludedSubnets.deleteMask(range);
  }

  void removeRange(const NetmaskGroup& group)
  {
    d_excludedSubnets.deleteMasks(group);
  }

  void excludeDomain(const DNSName& domain)
  {
    d_excludedDomains.add(domain);
  }

  std::string toString() const
  {
    std::stringstream result;

    result << "Query rate rule: " << d_queryRateRule.toString() << std::endl;
    result << "Response rate rule: " << d_respRateRule.toString() << std::endl;
    result << "SuffixMatch rule: " << d_suffixMatchRule.toString() << std::endl;
    result << "Response cache-miss ratio rule: " << d_respCacheMissRatioRule.toString() << std::endl;
    result << "RCode rules: " << std::endl;
    for (const auto& rule : d_rcodeRules) {
      result << "- " << RCode::to_s(rule.first) << ": " << rule.second.toString() << std::endl;
    }
    for (const auto& rule : d_rcodeRatioRules) {
      result << "- " << RCode::to_s(rule.first) << ": " << rule.second.toString() << std::endl;
    }
    result << "QType rules: " << std::endl;
    for (const auto& rule : d_qtypeRules) {
      result << "- " << QType(rule.first).toString() << ": " << rule.second.toString() << std::endl;
    }
    result << "Excluded Subnets: " << d_excludedSubnets.toString() << std::endl;
    result << "Excluded Domains: " << d_excludedDomains.toString() << std::endl;

    return result.str();
  }

  void setQuiet(bool quiet)
  {
    d_beQuiet = quiet;
  }

private:
  void applySMT(const struct timespec& now, StatNode& statNodeRoot);
  bool checkIfQueryTypeMatches(const Rings::Query& query);
  bool checkIfResponseCodeMatches(const Rings::Response& response);
  void addOrRefreshBlock(boost::optional<NetmaskTree<DynBlock, AddressAndPortRange>>& blocks, const struct timespec& now, const AddressAndPortRange& requestor, const DynBlockRule& rule, bool& updated, bool warning);
  void addOrRefreshBlockSMT(SuffixMatchTree<DynBlock>& blocks, const struct timespec& now, const DNSName& name, const DynBlockRule& rule, bool& updated);

  void addBlock(boost::optional<NetmaskTree<DynBlock, AddressAndPortRange>>& blocks, const struct timespec& now, const AddressAndPortRange& requestor, const DynBlockRule& rule, bool& updated)
  {
    addOrRefreshBlock(blocks, now, requestor, rule, updated, false);
  }

  void handleWarning(boost::optional<NetmaskTree<DynBlock, AddressAndPortRange>>& blocks, const struct timespec& now, const AddressAndPortRange& requestor, const DynBlockRule& rule, bool& updated)
  {
    addOrRefreshBlock(blocks, now, requestor, rule, updated, true);
  }

  bool hasQueryRules() const
  {
    return d_queryRateRule.isEnabled() || !d_qtypeRules.empty();
  }

  bool hasResponseRules() const
  {
    return d_respRateRule.isEnabled() || !d_rcodeRules.empty() || !d_rcodeRatioRules.empty() || d_respCacheMissRatioRule.isEnabled();
  }

  bool hasSuffixMatchRules() const
  {
    return d_suffixMatchRule.isEnabled();
  }

  bool hasRules() const
  {
    return hasQueryRules() || hasResponseRules();
  }

  void processQueryRules(counts_t& counts, const struct timespec& now);
  void processResponseRules(counts_t& counts, StatNode& root, const struct timespec& now);

  std::map<uint8_t, DynBlockRule> d_rcodeRules;
  std::map<uint8_t, DynBlockRatioRule> d_rcodeRatioRules;
  std::map<uint16_t, DynBlockRule> d_qtypeRules;
  DynBlockRule d_queryRateRule;
  DynBlockRule d_respRateRule;
  DynBlockRule d_suffixMatchRule;
  DynBlockCacheMissRatioRule d_respCacheMissRatioRule;
  NetmaskGroup d_excludedSubnets;
  SuffixMatchNode d_excludedDomains;
  smtVisitor_t d_smtVisitor;
  dnsdist_ffi_stat_node_visitor_t d_smtVisitorFFI;
  dnsdist_ffi_dynamic_block_inserted_hook d_newBlockHook;
  uint8_t d_v6Mask{128};
  uint8_t d_v4Mask{32};
  uint8_t d_portMask{0};
  bool d_beQuiet{false};
};

class DynBlockMaintenance
{
public:
  static void run();

  /* return the (cached) number of hits per second for the top offenders, averaged over 60s */
  static std::map<std::string, std::list<std::pair<AddressAndPortRange, unsigned int>>> getHitsForTopNetmasks();
  static std::map<std::string, std::list<std::pair<DNSName, unsigned int>>> getHitsForTopSuffixes();

  /* get the the top offenders based on the current value of the counters */
  static std::map<std::string, std::list<std::pair<AddressAndPortRange, unsigned int>>> getTopNetmasks(size_t topN);
  static std::map<std::string, std::list<std::pair<DNSName, unsigned int>>> getTopSuffixes(size_t topN);
  static void purgeExpired(const struct timespec& now);

  static time_t s_expiredDynBlocksPurgeInterval;

private:
  static void collectMetrics();
  static void generateMetrics();

  struct MetricsSnapshot
  {
    std::map<std::string, std::list<std::pair<AddressAndPortRange, unsigned int>>> nmgData;
    std::map<std::string, std::list<std::pair<DNSName, unsigned int>>> smtData;
  };

  struct Tops
  {
    std::map<std::string, std::list<std::pair<AddressAndPortRange, unsigned int>>> topNMGsByReason;
    std::map<std::string, std::list<std::pair<DNSName, unsigned int>>> topSMTsByReason;
  };

  static LockGuarded<Tops> s_tops;
  /* s_metricsData should only be accessed by the dynamic blocks maintenance thread so it does not need a lock */
  // need N+1 datapoints to be able to do the diff after a collection point has been reached
  static std::list<MetricsSnapshot> s_metricsData;
  static size_t s_topN;
};

namespace dnsdist::DynamicBlocks
{
bool addOrRefreshBlock(NetmaskTree<DynBlock, AddressAndPortRange>& blocks, const struct timespec& now, const AddressAndPortRange& requestor, const std::string& reason, unsigned int duration, DNSAction::Action action, bool warning, bool beQuiet);
bool addOrRefreshBlockSMT(SuffixMatchTree<DynBlock>& blocks, const struct timespec& now, const DNSName& name, const std::string& reason, unsigned int duration, DNSAction::Action action, bool beQuiet);
}
#endif /* DISABLE_DYNBLOCKS */