Merge pull request #52 from mgerst/feature/43-filter-autopwn

Limit AutoPWN by team, service name, or both.

flag_slurper/autopwn.py

@@ -22,8 +22,14 @@
 
 
 @click.group()
-def autopwn():
-    pass
+@click.pass_context
+def autopwn(ctx):
+    p = Project.get_instance()
+    if not p.enabled:
+        utils.report_error("AutoPWN commands require an active project")
+        exit(4)
+    p.connect_database()
+    ctx.obj = p
 
 
 def _pwn_service(limit_creds, service):
@@ -52,24 +58,30 @@ def _print_result(result, verbose):
 @autopwn.command()
 @click.option('-v', '--verbose', is_flag=True)
 @click.option('-P', '--parallel', is_flag=True, help="Async AutoPWN attack")
-@click.option('-N', '--processes', type=click.INT, default=None)
-@click.option('-c', '--limit-creds', type=click.STRING, multiple=True)
-def pwn(verbose, parallel, processes, limit_creds):
+@click.option('-N', '--processes', type=click.INT, default=None, help="How manny process to use for async AutoPWN")
+@click.option('-c', '--limit-creds', type=click.STRING, multiple=True, help="Limit the attack to the given creds")
+@click.option('-t', '--team', type=click.INT, default=None, help="Limit the attack to the given team")
+@click.option('-s', '--service', type=click.STRING, default=None, help="Limit the attack to the given service name")
+def pwn(verbose, parallel, processes, limit_creds, team, service):
     utils.report_status("Starting AutoPWN")
     p = Project.get_instance()
 
     if not processes:
         processes = os.cpu_count() + 1
 
-    if not p.enabled:
-        utils.report_error("AutoPwn requires a project be active")
-        return 1
-
     p.connect_database()
     utils.report_status("Loaded project from {}".format(p.base))
 
     services = models.Service.select()
 
+    if team:
+        utils.report_status('Limited to team {}'.format(team))
+        services = services.join(models.Team).where(models.Team.number == team)
+
+    if service:
+        utils.report_status('Limited to service {}'.format(service))
+        services = services.where(models.Service.service_name == service)
+
     if parallel:
         print("Using pool size: {}".format(processes))
         with Pool(processes=processes) as pool:
@@ -104,10 +116,6 @@ def display_service(cred: models.Credential):
 def generate(reconcile):
     p = Project.get_instance()
     p.connect_database()
-    if not p.enabled:
-        utils.report_error("Generate requires a project be active")
-        return 1
-
     teams = utils.get_teams()
 
     if reconcile:
@@ -136,11 +144,6 @@ def generate(reconcile):
 @autopwn.command()
 def results():
     p = Project.get_instance()
-
-    if not p.enabled:
-        utils.report_error("This command requires a project be active")
-        exit(3)
-
     p.connect_database()
 
     utils.report_status("Found the following flags")

tests/test_autopwn.py

@@ -0,0 +1,42 @@
+import pytest
+from click.testing import CliRunner
+
+from flag_slurper.cli import cli
+from flag_slurper.project import Project
+
+
+@pytest.fixture
+def pwn_project(create_project):
+    tmpdir = create_project("""
+    _version: "1.0"
+    project: Flag Slurper Test
+    base: {dir}/pwn-test
+    """)
+    p = Project.get_instance()
+    p.load(str(tmpdir.join('project.yml')))
+    return str(tmpdir.join('project.yml'))
+
+
+def test_autopwn_no_project():
+    p = Project.get_instance()
+    p.project_data = None
+    runner = CliRunner()
+    result = runner.invoke(cli, ['-np', 'autopwn', 'results'])
+    assert result.exit_code == 4
+    assert result.output == "[!] AutoPWN commands require an active project\n"
+
+
+def test_autopwn_pwn_limit_team(pwn_project, mocker, service):
+    runner = CliRunner()
+    pwn_service = mocker.patch('flag_slurper.autopwn._pwn_service')
+    result = runner.invoke(cli, ['autopwn', 'pwn', '-t', service.team.number])
+    assert result.exit_code == 0
+    pwn_service.assert_called_with((), service)
+
+
+def test_autopwn_pwn_limit_service(pwn_project, mocker):
+    runner = CliRunner()
+    pwn_service = mocker.patch('flag_slurper.autopwn._pwn_service')
+    result = runner.invoke(cli, ['autopwn', 'pwn', '-s', 'non-existant service'])
+    assert result.exit_code == 0
+    pwn_service.assert_not_called()