Merge pull request #31 from FCG-LLC/feature/aucote_status

Add REST API to the Aucote

api/handler.py

@@ -0,0 +1,70 @@
+"""
+Handler abstract class
+"""
+import hashlib
+
+from tornado.web import RequestHandler
+
+from aucote_cfg import cfg
+
+
+class Handler(RequestHandler):
+    """
+    Defines common properties for handler
+
+    """
+    def initialize(self, aucote):
+        """
+        Integrates Handlers with aucote
+
+        Args:
+            aucote (Aucote):
+
+        Returns:
+            None
+
+        """
+        self.aucote = aucote
+
+    @staticmethod
+    def auth(handler_class):
+        """
+        Handler for authorization
+
+        Args:
+            handler_class:
+
+        Returns:
+
+        """
+        def wrap_execute(handler_execute):
+            def require_auth(handler, kwargs):
+                auth_header = handler.request.headers.get('Authorization')
+
+                if auth_header is None or not auth_header.startswith('Bearer '):
+                    handler.set_status(401)
+                    handler._transforms = []
+                    handler.finish()
+                    return False
+
+                password = auth_header[7:]
+                hash = hashlib.sha512(password.encode()).hexdigest()
+                correct = cfg.get('service.api.password')
+
+                if hash != correct:
+                    handler.set_status(401)
+                    handler._transforms = []
+                    handler.finish()
+                    return False
+
+                return True
+
+            def _execute(self, transforms, *args, **kwargs):
+                if not require_auth(self, kwargs):
+                    return False
+                return handler_execute(self, transforms, *args, **kwargs)
+
+            return _execute
+
+        handler_class._execute = wrap_execute(handler_class._execute)
+        return handler_class

api/kill_handler.py

@@ -0,0 +1,25 @@
+"""
+Handler responsible for exit aucote application immediately
+
+"""
+import hashlib
+
+from api.handler import Handler
+from aucote_cfg import cfg
+
+@Handler.auth
+class KillHandler(Handler):
+    """
+    Kills aucote
+
+    """
+    def post(self):
+        """
+        Kills aucote. Require password POST argument
+
+        Returns:
+            None - kill application
+
+        """
+
+        self.aucote.kill()

api/main_handler.py

@@ -0,0 +1,164 @@
+"""
+Handler responsible for returning status of aucote
+
+"""
+import time
+
+from api.handler import Handler
+from aucote_cfg import cfg
+from scans.executor import Executor
+from tools.base import Tool
+from tools.common.port_task import PortTask
+from utils.task import Task
+
+
+class MainHandler(Handler):
+    """
+    Handler responsible for returning status of aucote
+
+    """
+    def get(self):
+        """
+        Handle get method and returns aucote status in JSON
+
+        Returns:
+            None - writes aucote status in JSON
+
+        """
+        self.write(self.aucote_status())
+
+    def aucote_status(self):
+        """
+        Get current status of aucote tasks
+
+        Returns:
+            dict
+
+        """
+        stats = self.thread_pool_status(self.aucote.thread_pool)
+        stats['scanner'] = self.scanning_status(self.aucote.scan_thread)
+        stats['meta'] = self.metadata()
+        return stats
+
+    @classmethod
+    def metadata(cls):
+        """
+        Meta of API request
+
+        Returns:
+            dict
+        """
+        return {
+            'timestamp': time.time()
+        }
+
+    @classmethod
+    def scanning_status(cls, scan_thread):
+        """
+        Information about scan
+
+        Args:
+            scan_thread (ScanThread):
+
+        Returns:
+            dict
+
+        """
+        return {
+            'nodes': [str(node.ip) for node in scan_thread.current_scan],
+            'scheduler': [cls.scheduler_task_status(task) for task in scan_thread.tasks],
+            'networks': cfg.get('service.scans.networks').cfg,
+            'ports': cfg.get('service.scans.ports'),
+            'previous_scan': scan_thread.previous_scan
+        }
+
+    @classmethod
+    def scheduler_task_status(cls, task):
+        """
+        Returns information about schedulers task
+
+        Args:
+            task : named tuple representing scheduler task
+
+        Returns:
+            dict
+
+        """
+        return {
+            'action': task.action.__name__,
+            'time': task.time
+        }
+
+    @classmethod
+    def thread_pool_status(cls, thread_pool):
+        """
+        Obtain status of thread pool
+
+        Args:
+            thread_pool (ThreadPool):
+
+        Returns:
+            dict
+
+        """
+        return_value = {'queue': [], 'threads': []}
+
+        for thread in thread_pool.threads:
+            return_value['threads'].append(cls.thread_pool_thread_status(thread))
+
+        for task in thread_pool.task_queue:
+            return_value['queue'].append(cls.task_status(task))
+
+        return_value['threads_limit'] = thread_pool.num_threads
+
+        return return_value
+
+    @classmethod
+    def thread_pool_thread_status(cls, thread):
+        """
+        Returns dict with info about thread
+
+        Args:
+            thread(Thread):
+
+        Returns:
+            dict
+
+        """
+        task = thread.task
+        if task is None:
+            return {}
+
+        return_value = cls.task_status(thread.task)
+
+        return return_value
+
+    @classmethod
+    def task_status(cls, task):
+        """
+        Returns information about task
+
+        Args:
+            task (Task):
+
+        Returns:
+            dict
+
+        """
+        return_value = {}
+
+        if isinstance(task, Task):
+            return_value['start_time'] = task.start_time
+            return_value['creation_time'] = task.creation_time
+            return_value['name'] = task.name
+
+        if isinstance(task, (Tool, PortTask)):
+            return_value['port'] = str(task.port)
+
+        if isinstance(task, Executor):
+            return_value['nodes'] = [str(node) for node in task.ports]
+
+        if isinstance(task, PortTask):
+            return_value['exploits'] = [exploit.name for exploit in task.current_exploits]
+
+        return return_value

aucote.py

@@ -11,13 +11,15 @@
 import fcntl
 
 import signal
+from threading import Lock
 
 from fixtures.exploits import Exploits
 from scans.executor_config import EXECUTOR_CONFIG
 from scans.task_mapper import TaskMapper
 from threads.scan_thread import ScanThread
 from threads.storage_thread import StorageThread
 from threads.watchdog_thread import WatchdogThread
+from threads.web_server_thread import WebServerThread
 from utils.exceptions import NmapUnsupported, TopdisConnectionException
 from utils.threads import ThreadPool
 from utils.kudu_queue import KuduQueue
@@ -96,6 +98,7 @@ class Aucote(object):
     """
 
     def __init__(self, exploits, kudu_queue, tools_config):
+        self._lock = Lock()
         self.exploits = exploits
         self._thread_pool = ThreadPool(cfg.get('service.scans.threads'))
         self._kudu_queue = kudu_queue
@@ -123,7 +126,8 @@ def storage(self):
             Storage
 
         """
-        return self._storage_thread
+        with self._lock:
+            return self._storage_thread
 
     @property
     def thread_pool(self):
@@ -154,14 +158,24 @@ def run_scan(self, as_service=True):
 
             self._scan_thread = ScanThread(aucote=self, as_service=as_service)
             self._scan_thread.start()
+
             self.thread_pool.start()
+            web_server = WebServerThread(self, cfg.get('service.api.v1.host'), cfg.get('service.api.v1.port'))
+            web_server.start()
 
-            self._scan_thread.join()
+            self.scan_thread.join()
             self.thread_pool.join()
 
+            web_server.stop()
+            web_server.join()
+
             self.thread_pool.stop()
-            self._storage_thread.stop()
-            self._storage_thread.join()
+            self.storage.stop()
+            self.storage.join()
+
+            self._scan_thread = None
+            self._watch_thread = None
+            self._storage_thread = None
 
         except TopdisConnectionException:
             log.exception("Exception while connecting to Topdis")
@@ -206,6 +220,18 @@ def signal_handler(self, sig, frame):
         log.error("Received signal %s at frame %s. Exiting.", sig, frame)
         self.kill()
 
+    @property
+    def scan_thread(self):
+        """
+        Scan thread
+
+        Returns:
+            ScanThread
+
+        """
+        with self._lock:
+            return self._scan_thread
+
     @property
     def unfinished_tasks(self):
         """

aucote_cfg.yaml.example

@@ -60,6 +60,13 @@ kuduworker:
 #       - <address/mask>
 #       - 127.0.0.1/32 - example ipv4 network
 #       - ::1/128 - example ipv6 network
+#   api:
+#     v1:
+#       host: <ip>  - IP on which web server should listen
+#       port: (int) - port on which web server should listen
+#     password: <hexstr> - SHA512 hash of password in hexadecimal format
+#                          Require for admin operations, e.g. killing aucote process.
+#                          Please do not use default password!
 
 service:
     scans:
@@ -75,6 +82,11 @@ service:
         # networks: # Change it to proper values and uncomment
         #    - 127.0.0.1/32
         #    - ::1/128
+    api:
+        v1:
+            host: 0.0.0.0
+            port: 1235
+        # password: ee26b0dd4af7e749aa1a8ee3c10ae9923f618980772e473f8819a5d4940e0db27ac185f8a0e1d5f84f88bc887fd67b143732c304cc5fa9ad8e6f57f50028a8ff # password: test
 
 # This section contains configuration related to topdis service
 # topdis:

requirements.txt

@@ -1,10 +1,10 @@
-pyaml
-argparse
-nanomsg
-netaddr
+pyaml==15.8.2
+nanomsg==1.0
+netaddr==0.7.18
 python-dateutil
-netifaces
-requests
-croniter
-pytz
-inotify
+netifaces==0.10.5
+requests==2.11.1
+croniter==0.3.13
+pytz==2016.10
+inotify==0.2.8
+tornado==4.4.2