text: /*
   * Copyright (C) 2004-2022 ZNC, see the NOTICE file for details.
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  #include <znc/znc.h>
  #include <znc/User.h>
  
  class CFailToBanMod : public CModule {
    public:
      MODCONSTRUCTOR(CFailToBanMod) {
          AddHelpCommand();
          AddCommand(
              "Timeout", t_d("[minutes]"),
              t_d("The number of minutes IPs are blocked after a failed login."),
              [=](const CString& sLine) { OnTimeoutCommand(sLine); });
          AddCommand("Attempts", t_d("[count]"),
                     t_d("The number of allowed failed login attempts."),
                     [=](const CString& sLine) { OnAttemptsCommand(sLine); });
          AddCommand("Ban", t_d("<hosts>"), t_d("Ban the specified hosts."),
                     [=](const CString& sLine) { OnBanCommand(sLine); });
          AddCommand("Unban", t_d("<hosts>"), t_d("Unban the specified hosts."),
                     [=](const CString& sLine) { OnUnbanCommand(sLine); });
          AddCommand("List", "", t_d("List banned hosts."),
                     [=](const CString& sLine) { OnListCommand(sLine); });
      }
      ~CFailToBanMod() override {}
  
      bool OnLoad(const CString& sArgs, CString& sMessage) override {
          CString sTimeout = sArgs.Token(0);
          CString sAttempts = sArgs.Token(1);
          unsigned int timeout = sTimeout.ToUInt();
  
          if (sAttempts.empty())
              m_uiAllowedFailed = 2;
          else
              m_uiAllowedFailed = sAttempts.ToUInt();
  
          if (sArgs.empty()) {
              timeout = 1;
          } else if (timeout == 0 || m_uiAllowedFailed == 0 ||
                     !sArgs.Token(2, true).empty()) {
              sMessage =
                  t_s("Invalid argument, must be the number of minutes IPs are "
                      "blocked after a failed login and can be followed by "
                      "number of allowed failed login attempts");
              return false;
          }
  
          // SetTTL() wants milliseconds
          m_Cache.SetTTL(timeout * 60 * 1000);
  
          return true;
      }
  
      void OnPostRehash() override { m_Cache.Clear(); }
  
      void Add(const CString& sHost, unsigned int count) {
          m_Cache.AddItem(sHost, count, m_Cache.GetTTL());
      }
  
      bool Remove(const CString& sHost) { return m_Cache.RemItem(sHost); }
  
      void OnTimeoutCommand(const CString& sCommand) {
          if (!GetUser()->IsAdmin()) {
              PutModule(t_s("Access denied"));
              return;
          }
  
          CString sArg = sCommand.Token(1);
  
          if (!sArg.empty()) {
              unsigned int uTimeout = sArg.ToUInt();
              if (uTimeout == 0) {
                  PutModule(t_s("Usage: Timeout [minutes]"));
              } else {
                  m_Cache.SetTTL(uTimeout * 60 * 1000);
                  SetArgs(CString(m_Cache.GetTTL() / 60 / 1000) + " " +
                          CString(m_uiAllowedFailed));
                  PutModule(t_f("Timeout: {1} min")(uTimeout));
              }
          } else {
              PutModule(t_f("Timeout: {1} min")(m_Cache.GetTTL() / 60 / 1000));
          }
      }
  
      void OnAttemptsCommand(const CString& sCommand) {
          if (!GetUser()->IsAdmin()) {
              PutModule(t_s("Access denied"));
              return;
          }
  
          CString sArg = sCommand.Token(1);
  
          if (!sArg.empty()) {
              unsigned int uiAttempts = sArg.ToUInt();
              if (uiAttempts == 0) {
                  PutModule(t_s("Usage: Attempts [count]"));
              } else {
                  m_uiAllowedFailed = uiAttempts;
                  SetArgs(CString(m_Cache.GetTTL() / 60 / 1000) + " " +
                          CString(m_uiAllowedFailed));
                  PutModule(t_f("Attempts: {1}")(uiAttempts));
              }
          } else {
              PutModule(t_f("Attempts: {1}")(m_uiAllowedFailed));
          }
      }
  
      void OnBanCommand(const CString& sCommand) {
          if (!GetUser()->IsAdmin()) {
              PutModule(t_s("Access denied"));
              return;
          }
  
          CString sHosts = sCommand.Token(1, true);
  
          if (sHosts.empty()) {
              PutStatus(t_s("Usage: Ban <hosts>"));
              return;
          }
  
          VCString vsHosts;
          sHosts.Replace(",", " ");
          sHosts.Split(" ", vsHosts, false, "", "", true, true);
  
          for (const CString& sHost : vsHosts) {
              Add(sHost, 0);
              PutModule(t_f("Banned: {1}")(sHost));
          }
      }
  
      void OnUnbanCommand(const CString& sCommand) {
          if (!GetUser()->IsAdmin()) {
              PutModule(t_s("Access denied"));
              return;
          }
  
          CString sHosts = sCommand.Token(1, true);
  
          if (sHosts.empty()) {
              PutStatus(t_s("Usage: Unban <hosts>"));
              return;
          }
  
          VCString vsHosts;
          sHosts.Replace(",", " ");
          sHosts.Split(" ", vsHosts, false, "", "", true, true);
  
          for (const CString& sHost : vsHosts) {
              if (Remove(sHost)) {
                  PutModule(t_f("Unbanned: {1}")(sHost));
              } else {
                  PutModule(t_f("Ignored: {1}")(sHost));
              }
          }
      }
  
      void OnListCommand(const CString& sCommand) {
          if (!GetUser()->IsAdmin()) {
              PutModule(t_s("Access denied"));
              return;
          }
  
          CTable Table;
          Table.AddColumn(t_s("Host", "list"));
          Table.AddColumn(t_s("Attempts", "list"));
          Table.SetStyle(CTable::ListStyle);
  
          for (const auto& it : m_Cache.GetItems()) {
              Table.AddRow();
              Table.SetCell(t_s("Host", "list"), it.first);
              Table.SetCell(t_s("Attempts", "list"), CString(it.second));
          }
  
          if (Table.empty()) {
              PutModule(t_s("No bans", "list"));
          } else {
              PutModule(Table);
          }
      }
  
      void OnClientConnect(CZNCSock* pClient, const CString& sHost,
                           unsigned short uPort) override {
          unsigned int* pCount = m_Cache.GetItem(sHost);
          if (sHost.empty() || pCount == nullptr || *pCount < m_uiAllowedFailed) {
              return;
          }
  
          // refresh their ban
          Add(sHost, *pCount);
  
          pClient->Write(
              "ERROR :Closing link [Please try again later - reconnecting too "
              "fast]\r\n");
          pClient->Close(Csock::CLT_AFTERWRITE);
      }
  
      void OnFailedLogin(const CString& sUsername,
                         const CString& sRemoteIP) override {
          unsigned int* pCount = m_Cache.GetItem(sRemoteIP);
          if (pCount)
              Add(sRemoteIP, *pCount + 1);
          else
              Add(sRemoteIP, 1);
      }
  
      EModRet OnLoginAttempt(std::shared_ptr<CAuthBase> Auth) override {
          // e.g. webadmin ends up here
          const CString& sRemoteIP = Auth->GetRemoteIP();
  
          if (sRemoteIP.empty()) return CONTINUE;
  
          unsigned int* pCount = m_Cache.GetItem(sRemoteIP);
          if (pCount && *pCount >= m_uiAllowedFailed) {
              // OnFailedLogin() will refresh their ban
              Auth->RefuseLogin("Please try again later - reconnecting too fast");
              return HALT;
          }
  
          return CONTINUE;
      }
  
    private:
      TCacheMap<CString, unsigned int> m_Cache;
      unsigned int m_uiAllowedFailed{};
  };
  
  template <>
  void TModInfo<CFailToBanMod>(CModInfo& Info) {
      Info.SetWikiPage("fail2ban");
      Info.SetHasArgs(true);
      Info.SetArgsHelpText(
          Info.t_s("You might enter the time in minutes for the IP banning and "
                   "the number of failed logins before any action is taken."));
  }
  
  GLOBALMODULEDEFS(CFailToBanMod,
                   t_s("Block IPs for some time after a failed login."))
  