Repository: mgerst/flag-slurper

Issue:
Add delay to AutoPWN to avoid fail2ban
Fail2ban's default configuration bans an IP if it fails to login 3 times in a 10 minutes window. The default credential list, *in theory*, should not trip this since it only makes two attempts (`root:cdc` and `cdc:cdc`).

AutoPWN should either add a delay in between each attempt or add a delay after `maxretry - 1` attempts. Either way it should be configurable.

Discussion:

Patch:
diff --git a/flag_slurper/autolib/governor.py b/flag_slurper/autolib/governor.py
new file mode 100644
index 0000000..5ace3ae
--- /dev/null
+++ b/flag_slurper/autolib/governor.py
@@ -0,0 +1,64 @@
+import logging
+import socket
+import time
+from collections import defaultdict
+from datetime import datetime, timedelta
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class Governor:
+    """
+    The governor tracks the number of connections over time
+    to a given ip address. This is used for fail2ban evasion.
+
+    This can be configured in the flagrc.
+
+    .. warning::
+
+       The governor only tracks withing it's worker process, for a single run.
+       It **DOES NOT** persist across invocations.
+
+    * ``autopwn->governor`` (default false) whether the governor is enabled.
+    * ``autopwn->delay`` (default 5m) how long to wait after limiting.
+    * ``autopwn->window`` (default 30m) the window to track attempts in.
+    * ``autopwn->times`` (default 3) max attempts inside of the window.
+    """
+    instance = None
+
+    def __init__(self, enabled: bool, delay: int = 5 * 60, window: int = 30 * 60, times: int = 3):
+        self.enabled = enabled
+        self.delay = delay
+        self.window = window
+        self.times = times
+
+        # Dict[IP, timestamp]
+        self.limits = defaultdict(list)
+
+    @classmethod
+    def get_instance(cls, enabled: Optional[bool] = False, delay: Optional[int] = 5 * 60,
+                     window: Optional[int] = 30 * 60, times: Optional[int] = 3) -> 'Governor':
+        if not Governor.instance:
+            Governor.instance = cls(enabled, delay, window, times)
+        return Governor.instance
+
+    def filter(self, ipaddr: str):
+        def _filter(x: datetime):
+            return x > datetime.now() - timedelta(seconds=self.window)
+
+        self.limits[ipaddr] = list(filter(_filter, self.limits[ipaddr]))
+
+    def attempt(self, ipaddr: str):
+        if not self.enabled:
+            return
+
+        self.filter(ipaddr)
+        self.limits[ipaddr].append(datetime.now())
+        if len(self.limits[ipaddr]) > self.times:
+            logger.info("Attempting %d second delay for %s", self.delay, ipaddr)
+            time.sleep(self.delay)
+
+    @staticmethod
+    def resolve_url(url: str) -> str:
+        return socket.gethostbyname(url)
diff --git a/flag_slurper/autolib/protocols.py b/flag_slurper/autolib/protocols.py
index 8abf1d9..a655a8e 100644
--- a/flag_slurper/autolib/protocols.py
+++ b/flag_slurper/autolib/protocols.py
@@ -5,6 +5,7 @@ from typing import Tuple
 import paramiko
 
 from flag_slurper.autolib.exploit import get_file_contents, get_system_info
+from flag_slurper.autolib.governor import Governor
 from .exploit import find_flags, FlagConf, can_sudo
 from .models import Service, CredentialBag, Credential, Flag, CaptureNote
 from .post import post_ssh
@@ -26,6 +27,10 @@ def pwn_ssh(url: str, port: int, service: Service, flag_conf: FlagConf) -> Tuple
     working = set()
     credentials = CredentialBag.select()
     for credential in credentials:
+        # Govern if necessary (and enabled)
+        gov = Governor.get_instance()
+        gov.attempt(gov.resolve_url(url))
+
         try:
             cred = credential.credentials.where(Credential.service == service).get()
         except Credential.DoesNotExist:
diff --git a/flag_slurper/autopwn.py b/flag_slurper/autopwn.py
index b0a625d..4d77997 100644
--- a/flag_slurper/autopwn.py
+++ b/flag_slurper/autopwn.py
@@ -4,6 +4,7 @@ from multiprocessing import Pool
 
 import click
 
+from flag_slurper.autolib.governor import Governor
 from flag_slurper.autolib.models import SUDO_FLAG
 from . import utils, autolib
 from .autolib import models
diff --git a/flag_slurper/config.py b/flag_slurper/config.py
index 2f43948..09c969a 100644
--- a/flag_slurper/config.py
+++ b/flag_slurper/config.py
@@ -7,6 +7,8 @@ import requests
 from click import echo, prompt
 from jinja2 import Environment
 
+from flag_slurper.autolib.governor import Governor
+
 ROOT = Path(__file__).parent
 
 
@@ -31,9 +33,19 @@ class Config(ConfigParser):
             conffiles.append(extra)
         conf.read(conffiles)
 
+        conf.configure_governor()
+
         cls.instance = conf
         return conf
 
+    def configure_governor(self):
+        from .utils import parse_duration
+        gov = Governor.get_instance()
+        gov.enabled = self.getboolean('autopwn', 'governor')
+        gov.delay = parse_duration(self['autopwn']['delay'])
+        gov.window = parse_duration(self['autopwn']['window'])
+        gov.times = self.getint('autopwn', 'times')
+
     @staticmethod
     def get_instance():
         if not Config.instance:
diff --git a/flag_slurper/default.ini b/flag_slurper/default.ini
index 5ab9cbd..72ae3f2 100644
--- a/flag_slurper/default.ini
+++ b/flag_slurper/default.ini
@@ -5,3 +5,9 @@ api_version=v1
 
 [database]
 url=sqlite:///{{ project }}/db.sqlite3
+
+[autopwn]
+governor=false
+delay=5m
+window=30m
+times=3
diff --git a/flag_slurper/utils.py b/flag_slurper/utils.py
index d44fa83..b06dd61 100644
--- a/flag_slurper/utils.py
+++ b/flag_slurper/utils.py
@@ -2,6 +2,7 @@ import os
 import shutil
 import zipfile
 from typing import Tuple, Dict, Union, Callable, Optional
+import string
 
 import click
 import requests
@@ -167,3 +168,43 @@ def conditional_page(output: str, size: int):
         click.echo_via_pager(output)
     else:
         click.echo(output)
+
+
+def parse_duration(duration: str) -> int:
+    """
+    Parses a duration of the form 10m or 5s and converts
+    them into seconds.
+
+    >>> parse_duration('5s')
+    ... 5
+    >>> parse_duration('10m')
+    ... 600
+
+    .. note::
+
+       Only one unit is supported. For example 1h30m is not supported, but 90m is.
+
+    Supported suffixes:
+
+    - ``h`` hours
+    - ``m`` minutes
+    - ``s`` seconds
+
+    An omitted suffix will be interpreted as seconds.
+
+    :param duration: The duration in string form
+    :return: The duration in seconds
+    """
+    if duration[-1] in string.digits:
+        return int(duration)
+
+    if duration[-1].lower() == 'h':
+        return int(duration[:-1]) * 60 * 60
+
+    if duration[-1].lower() == 'm':
+        return int(duration[:-1]) * 60
+
+    if duration[-1].lower() == 's':
+        return int(duration[:-1])
+
+    raise ValueError("Unable to parse {duration} as a duration".format(duration=duration))


Test patch:
diff --git a/tests/autolib/test_governor.py b/tests/autolib/test_governor.py
new file mode 100644
index 0000000..f9812b5
--- /dev/null
+++ b/tests/autolib/test_governor.py
@@ -0,0 +1,64 @@
+from datetime import datetime, timedelta
+
+from flag_slurper.autolib.governor import Governor
+
+
+def test_governor_singleton():
+    Governor.instance = None
+    Governor.get_instance(True)
+    gov = Governor.get_instance()
+    assert gov.enabled
+    assert gov.delay == 5 * 60
+    assert gov.window == 30 * 60
+    assert gov.times == 3
+
+
+def test_governor_tracks(mock):
+    Governor.instance = None
+    gov = Governor.get_instance(True)
+    sleep = mock.patch('flag_slurper.autolib.governor.time.sleep')
+    gov.attempt('192.168.1.113')
+    assert not sleep.called
+
+
+def test_governor_filter():
+    Governor.instance = None
+    gov = Governor.get_instance(True)
+    gov.limits['192.168.1.113'] = [datetime.now() - timedelta(minutes=3),
+                                   datetime.now() - timedelta(minutes=29),
+                                   datetime.now() - timedelta(minutes=1),
+                                   datetime.now() - timedelta(minutes=35)]
+    gov.filter('192.168.1.113')
+    assert len(gov.limits['192.168.1.113']) == 3
+
+
+def test_governor_limits(mock):
+    Governor.instance = None
+    gov = Governor.get_instance(True)
+    gov.limits['192.168.1.113'] = [datetime.now() - timedelta(minutes=3),
+                                   datetime.now() - timedelta(minutes=29),
+                                   datetime.now() - timedelta(minutes=1),
+                                   datetime.now() - timedelta(minutes=35)]
+    sleep = mock.patch('flag_slurper.autolib.governor.time.sleep')
+    gov.attempt('192.168.1.113')
+    assert sleep.called
+
+
+def test_governer_doesnt_limit_when_disabled(mock):
+    Governor.instance = None
+    gov = Governor.get_instance(False)
+    gov.limits['192.168.1.113'] = [datetime.now() - timedelta(minutes=3),
+                                   datetime.now() - timedelta(minutes=29),
+                                   datetime.now() - timedelta(minutes=1),
+                                   datetime.now() - timedelta(minutes=35)]
+    sleep = mock.patch('flag_slurper.autolib.governor.time.sleep')
+    gov.attempt('192.168.1.113')
+    assert not sleep.called
+
+
+def test_governor_resolve(mock):
+    addrhost = mock.patch('flag_slurper.autolib.governor.socket.gethostbyname')
+    addrhost.return_value = '192.168.1.113'
+    ipaddr = Governor.resolve_url('example.com')
+    assert addrhost.called_with('example.com')
+    assert ipaddr == '192.168.1.113'
diff --git a/tests/test_config.py b/tests/test_config.py
index f142137..565e575 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -4,6 +4,7 @@ from unittest import mock
 
 from click.testing import CliRunner
 
+from flag_slurper.autolib.governor import Governor
 from flag_slurper.cli import cli
 from flag_slurper.config import Config
 from flag_slurper.models import User
@@ -97,3 +98,13 @@ def test_config_login(tmpdir):
     c.read_file(tmpdir.join('flagrc').open('r'))
     assert 'iscore' in c and 'api_token' in c['iscore']
     assert c['iscore']['api_token'] == 'abcdef'
+
+
+def test_config_governor(tmpdir):
+    rc = tmpdir.join("flagrc.ini")
+    rc.write("[autopwn]\ngovernor=yes\ndelay=6m\n")
+
+    Config.load(str(rc), noflagrc=True)
+    gov = Governor.get_instance()
+    assert gov.enabled
+    assert gov.delay == 6 * 60
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 14bd39a..0dd9ed6 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -168,3 +168,25 @@ def test_get_teams(config):
 def test_parse_creds(given, expected):
     result = utils.parse_creds(given)
     assert result == expected
+
+
+@pytest.mark.parametrize('given,expected', (
+    ('3', 3),
+    ('5s', 5),
+    ('10m', 10 * 60),
+    ('3h', 3 * 60 * 60),
+))
+def test_parse_duration(given, expected):
+    result = utils.parse_duration(given)
+    assert result == expected
+
+
+@pytest.mark.parametrize('given', (
+    'abc',
+    '30y',
+    '2d',
+))
+def test_parse_duration_invalid(given):
+    with pytest.raises(ValueError, match=r'Unable to parse \w+ as a duration$'):
+        utils.parse_duration(given)
+