Repository: ameserole/Akeso

Issue:
Add Support for Multiple Service and Exploit Checks
Currently only allows for one Service and Exploit check at a time. 

Comments:

Patch:
diff --git a/Akeso/AttackWorkers.py b/Akeso/AttackWorkers.py
index 97afaaf..87def44 100644
--- a/Akeso/AttackWorkers.py
+++ b/Akeso/AttackWorkers.py
@@ -9,14 +9,6 @@ from ServiceManager import ServiceInfo
 logger = structlog.get_logger()
 
 
-def challenge_mapper(challenge):
-    return {
-        'maze': ('maze', 'mazeAttack', 'maze', 31337),
-        'SQL': ('sqlisimple', 'SQLi', 'SQLiSimple', 80),
-        'shell': ('shell', 'shellAttack', 'shell', 4001),
-    }[challenge]
-
-
 def attackCallback(ch, method, properties, body):
     """Pull service off of attack queue and run selected attack against it"""
 
@@ -26,7 +18,7 @@ def attackCallback(ch, method, properties, body):
     resultChannel.queue_declare(queue='resultQueue', durable=True)
 
     body = json.loads(body)
-    info = challenge_mapper(body['chal'])
+    info = config.challenge_mapper(body['chal'])
 
     if 'serviceName' in body:
         serviceName = body['serviceName']
@@ -37,8 +29,8 @@ def attackCallback(ch, method, properties, body):
         'serviceName': serviceName,
         'imageName': info[0],
         'userInfo': body['userInfo'],
-        'exploitModule': info[1],
-        'serviceCheckName': info[2],
+        'exploitModules': info[1],
+        'serviceCheckNames': info[2],
         'serviceHost': body['serviceHost'],
         'servicePort': info[3]
     }
@@ -51,48 +43,69 @@ def attackCallback(ch, method, properties, body):
 
     log = logger.bind(service=service.__dict__)
     userMsg = "Starting Attack on {} {}\n".format(service.imageName, service.userInfo)
-
-    # Get the Service module for this service and check that it is running correctly
-    serviceCheckModuleName = 'Services.' + service.serviceCheckName + '.' + service.serviceCheckName
-    serviceModule = importlib.import_module(serviceCheckModuleName, package=None)
-    serviceCheckObject = serviceModule.ServiceCheck(service)
-
-    if serviceCheckObject.checkService():
-        log.info('attackCallback', msg="Service Check Succeeded")
-        userMsg = "Service Check Succeeded"
-    else:
-        log.info('attackCallback', msg="Service Check Failed")
-        userMsg = "Service Check Failed"
-        resultChannel.basic_publish(exchange='resultX',
-                                    routing_key=str(service.userInfo),
-                                    body=json.dumps({'msg': userMsg, 'service': service.__dict__}))
-        ch.basic_ack(delivery_tag=method.delivery_tag)
-        return -1
-
-    # If the service is running correctly grab the selected exploit module and run it against the current service
-    exploitModuleName = 'Exploits.' + service.exploitModule
-    exploitModule = importlib.import_module(exploitModuleName, package=None)
-    exploitObject = exploitModule.Exploit(service)
-    exploitObject.exploit()
-
-    exploitSuccess = exploitObject.exploitSuccess()
-
-    if exploitSuccess:
-        userMsg = "Your Code/Config was exploited."
-        log.info("attackCallback", msg="Exploit Success")
-        resultChannel.basic_publish(exchange='resultX',
-                                    routing_key=str(service.userInfo),
-                                    body=json.dumps({'msg': userMsg, 'service': service.__dict__}))
-
-        ch.basic_ack(delivery_tag=method.delivery_tag)
-        return -1
-
-    else:
-        userMsg = "Attack Failed"
-        log.info("attackCallback", msg=userMsg)
+    for serviceCheckName in service.serviceCheckNames:
+        # Get the Service module for this service and check that it is running correctly
+        serviceCheckModuleName = 'Services.' + serviceCheckName + '.' + serviceCheckName
+        serviceModule = importlib.import_module(serviceCheckModuleName, package=None)
+        serviceCheckObject = serviceModule.ServiceCheck(service)
+
+        if serviceCheckObject.checkService():
+            log.info('attackCallback', msg="Service Check Succeeded")
+            userMsg = "Service Check Succeeded"
+        else:
+            log.info('attackCallback', msg="Service Check Failed")
+            userMsg = "Service Check Failed"
+            resultChannel.basic_publish(exchange='resultX',
+                                        routing_key=str(service.userInfo),
+                                        body=json.dumps({'msg': userMsg, 'service': service.__dict__}))
+            ch.basic_ack(delivery_tag=method.delivery_tag)
+            return -1
+
+    for exploitModule in service.exploitModules:
+        # If the service is running correctly grab the selected exploit module and run it against the current service
+        exploitModuleName = 'Exploits.' + exploitModule
+        exploitModule = importlib.import_module(exploitModuleName, package=None)
+        exploitObject = exploitModule.Exploit(service)
+        exploitObject.exploit()
+
+        exploitSuccess = exploitObject.exploitSuccess()
+
+        if exploitSuccess:
+            userMsg = "Your Code/Config was exploited."
+            log.info("attackCallback", msg="Exploit Success")
+            resultChannel.basic_publish(exchange='resultX',
+                                        routing_key=str(service.userInfo),
+                                        body=json.dumps({'msg': userMsg, 'service': service.__dict__}))
+
+            ch.basic_ack(delivery_tag=method.delivery_tag)
+            return -1
+
+        else:
+            userMsg = "Attack Failed"
+            log.info("attackCallback", msg=userMsg)
 
     # Check to see if the service is still up after the exploit was run
-    checkService = serviceCheckObject.checkService()
+    # checkService = serviceCheckObject.checkService()
+
+    checkService = False
+    for serviceCheckName in service.serviceCheckNames:
+        # Get the Service module for this service and check that it is running correctly
+        serviceCheckModuleName = 'Services.' + serviceCheckName + '.' + serviceCheckName
+        serviceModule = importlib.import_module(serviceCheckModuleName, package=None)
+        serviceCheckObject = serviceModule.ServiceCheck(service)
+
+        checkService = serviceCheckObject.checkService()
+        if checkService:
+            log.info('attackCallback', msg="Service Check Succeeded")
+            userMsg = "Service Check Succeeded"
+        else:
+            log.info('attackCallback', msg="Service Check Failed After Attack")
+            userMsg = "Service Check Failed"
+            resultChannel.basic_publish(exchange='resultX',
+                                        routing_key=str(service.userInfo),
+                                        body=json.dumps({'msg': userMsg, 'service': service.__dict__}))
+            ch.basic_ack(delivery_tag=method.delivery_tag)
+            return -1
 
     # If Service is still up and exploit did not work return the flag to the user
     if not exploitSuccess and checkService:
@@ -102,17 +115,6 @@ def attackCallback(ch, method, properties, body):
                                     routing_key=str(service.userInfo),
                                     body=json.dumps({'msg': userMsg, 'service': service.__dict__}))
 
-    # No flag for you :(
-    elif not exploitSuccess and not checkService:
-        log.info('attackCallback', msg="Service Check Failed After Attack")
-        userMsg = "Service Check Failed After Attack"
-        resultChannel.basic_publish(exchange='resultX',
-                                    routing_key=str(service.userInfo),
-                                    body=json.dumps({'msg': userMsg, 'service': service.__dict__}))
-
-        ch.basic_ack(delivery_tag=method.delivery_tag)
-        return -1
-
     ch.basic_ack(delivery_tag=method.delivery_tag)
     return 1
 
diff --git a/Akeso/ServiceManager.py b/Akeso/ServiceManager.py
index 503e486..6557d8d 100644
--- a/Akeso/ServiceManager.py
+++ b/Akeso/ServiceManager.py
@@ -10,6 +10,6 @@ class ServiceInfo(object):
         self.imageName = info['imageName']
         self.serviceHost = info['serviceHost']
         self.servicePort = info['servicePort']
-        self.exploitModule = info['exploitModule']
-        self.serviceCheckName = info['serviceCheckName']
+        self.exploitModules = info['exploitModules']
+        self.serviceCheckNames = info['serviceCheckNames']
         self.userInfo = info['userInfo']
diff --git a/Akeso/config.py b/Akeso/config.py
index 2e6e884..a1718c7 100644
--- a/Akeso/config.py
+++ b/Akeso/config.py
@@ -8,3 +8,12 @@ SERVICE_PATH = os.path.join(os.getcwd(), 'Services/')
 
 # Address of the RabbitMQ server
 RABBITMQ_SERVER = '172.17.0.2'
+
+
+def challenge_mapper(challenge):
+    return {
+        'maze': ('maze', ['mazeAttack'], ['maze'], 31337),
+        'SQL': ('sqlisimple', ['SQLi'], ['SQLiSimple'], 80),
+        'shell': ('shell', ['shellAttack'], ['shell'], 4001),
+        'nginx': ('nginx', ['DirectoryTraversal'], ['ApacheDirectoryTraversal'], 80)
+    }[challenge]


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 3b56e0c..b11e9b4 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -10,8 +10,8 @@ def fake_service():
         'volumeLocation': 'fakeVolume',
         'serviceHost': '127.0.0.1',
         'servicePort': 80,
-        'exploitModule': 'fakeExploit',
-        'serviceCheckName': 'fakeCheck',
+        'exploitModules': ['fakeExploit'],
+        'serviceCheckNames': ['fakeCheck'],
         'userInfo': '1'}
 
     return ServiceInfo(fakeInfo)
diff --git a/tests/test_ServiceManager.py b/tests/test_ServiceManager.py
index c714d7b..28c991f 100644
--- a/tests/test_ServiceManager.py
+++ b/tests/test_ServiceManager.py
@@ -7,8 +7,8 @@ def test_ServiceInfo():
         'imageName': 'fakeImage',
         'serviceHost': '127.0.0.1',
         'servicePort': 80,
-        'exploitModule': 'fakeExploit',
-        'serviceCheckName': 'fakeCheck',
+        'exploitModules': ['fakeExploit'],
+        'serviceCheckNames': ['fakeCheck'],
         'userInfo': 'fakeInfo'}
 
     service = ServiceInfo(fakeInfo)
diff --git a/tests/test_config.py b/tests/test_config.py
new file mode 100644
index 0000000..6434cad
--- /dev/null
+++ b/tests/test_config.py
@@ -0,0 +1,13 @@
+import pytest
+from Akeso.config import challenge_mapper
+
+
+@pytest.mark.parametrize("challenge, expected_ret", [
+    ('maze', ('maze', ['mazeAttack'], ['maze'], 31337)),
+    ('SQL', ('sqlisimple', ['SQLi'], ['SQLiSimple'], 80)),
+    ('shell', ('shell', ['shellAttack'], ['shell'], 4001)),
+    ('nginx', ('nginx', ['DirectoryTraversal'], ['ApacheDirectoryTraversal'], 80))
+])
+def test_challenge_mapper(challenge, expected_ret):
+    ret = challenge_mapper(challenge)
+    assert ret == expected_ret