diff --git a/modules/fail2ban.cpp b/modules/fail2ban.cpp index a69fac7e..3d20ea62 100644 --- a/modules/fail2ban.cpp +++ b/modules/fail2ban.cpp @@ -14,13 +14,21 @@ public: virtual ~CFailToBanMod() {} virtual bool OnLoad(const CString& sArgs, CString& sMessage) { - unsigned int timeout = sArgs.ToUInt(); + 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 && sArgs != "0") { + } else if (timeout == 0 || m_uiAllowedFailed == 0 || !sArgs.Token(2, true).empty()) { sMessage = "Invalid argument, must be the number of minutes " - "IPs are blocked after a failed login"; + "IPs are blocked after a failed login and can be " + "followed by number of allowed failed login attempts"; return false; } @@ -30,6 +38,10 @@ public: return true; } + void Add(const CString& sHost, unsigned int count) { + m_Cache.AddItem(sHost, count, m_Cache.GetTTL()); + } + virtual void OnModCommand(const CString& sCommand) { PutModule("This module can only be configured through its arguments."); PutModule("The module argument is the number of minutes an IP"); @@ -37,19 +49,24 @@ public: } virtual void OnClientConnect(CClient* pClient, const CString& sHost, unsigned short uPort) { - if (sHost.empty() || !m_Cache.HasItem(sHost)) { + unsigned int *pCount = m_Cache.GetItem(sHost); + if (sHost.empty() || pCount == NULL || *pCount < m_uiAllowedFailed) { return; } // refresh their ban - m_Cache.AddItem(sHost); + Add(sHost, *pCount); pClient->PutClient("ERROR :Closing link [Please try again later - reconnecting too fast]"); pClient->Close(Csock::CLT_AFTERWRITE); } virtual void OnFailedLogin(const CString& sUsername, const CString& sRemoteIP) { - m_Cache.AddItem(sRemoteIP); + unsigned int *pCount = m_Cache.GetItem(sRemoteIP); + if (pCount) + Add(sRemoteIP, *pCount + 1); + else + Add(sRemoteIP, 1); } virtual EModRet OnLoginAttempt(CSmartPtr Auth) { @@ -59,8 +76,9 @@ public: if (sRemoteIP.empty()) return CONTINUE; - if (m_Cache.HasItem(sRemoteIP)) { - // OnFailedLoginAttempt() will refresh their ban + 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; } @@ -69,7 +87,8 @@ public: } private: - TCacheMap m_Cache; + TCacheMap m_Cache; + unsigned int m_uiAllowedFailed; }; GLOBALMODULEDEFS(CFailToBanMod, "Block IPs for some time after a failed login")