Merge pull request #167 from pemontto/feature-nessus-stream

Feature nessus stream

.gitignore

@@ -4,6 +4,7 @@ logs/
 elk6/vulnwhisperer.ini
 resources/elk6/vulnwhisperer.ini
 configs/frameworks_example.ini
+tests/data
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

.gitmodules

@@ -1,3 +1,3 @@
-[submodule "test"]
+[submodule "tests/data"]
-	path = test
+	path = tests/data
 	url = https://github.com/HASecuritySolutions/VulnWhisperer-tests

.travis.yml

@@ -3,7 +3,8 @@ language: python
 cache: pip
 python:
     - 2.7
-
+env:
+    - TEST_PATH=tests/data
 #    - 3.6
 #matrix:
 #    allow_failures:
@@ -20,21 +21,22 @@ before_script:
 script:
     - python setup.py install
     # Test successful scan download and parsing
-    - vuln_whisperer -c configs/test.ini --mock --mock_dir test
     - rm -rf /tmp/VulnWhisperer
+    - vuln_whisperer -c configs/test.ini --mock --mock_dir ${TEST_PATH}
     # Test one failed scan
-    - rm -f test/nessus/GET_scans_exports_164_download
-    - vuln_whisperer -c configs/test.ini --mock --mock_dir test; [[ $? -eq 1 ]]
     - rm -rf /tmp/VulnWhisperer
+    - rm -f ${TEST_PATH}/nessus/GET_scans_exports_164_download
+    - vuln_whisperer -c configs/test.ini --mock --mock_dir ${TEST_PATH}; [[ $? -eq 1 ]]
     # Test two failed scans
-    - rm -f test/qualys_vuln/scan_1553941061.87241
-    - vuln_whisperer -c configs/test.ini --mock --mock_dir test; [[ $? -eq 2 ]]
     - rm -rf /tmp/VulnWhisperer
+    - rm -f ${TEST_PATH}/qualys_vuln/scan_1553941061.87241
+    - vuln_whisperer -c configs/test.ini --mock --mock_dir ${TEST_PATH}; [[ $? -eq 2 ]]
     # Test only nessus
-    - vuln_whisperer -c configs/test.ini -s nessus --mock --mock_dir test; [[ $? -eq 1 ]]
     - rm -rf /tmp/VulnWhisperer
+    - vuln_whisperer -c configs/test.ini -s nessus --mock --mock_dir ${TEST_PATH}; [[ $? -eq 1 ]]
     # Test only qualy_vuln
-    - vuln_whisperer -c configs/test.ini -s qualys_vuln --mock --mock_dir test; [[ $? -eq 1 ]]
+    - rm -rf /tmp/VulnWhisperer
+    - vuln_whisperer -c configs/test.ini -s qualys_vuln --mock --mock_dir ${TEST_PATH}; [[ $? -eq 1 ]]
 notifications:
     on_success: change
     on_failure: change  # `always` will be the setting once code changes slow down

bin/vuln_whisperer

@@ -37,10 +37,13 @@ def main():
                         help='The NESSUS username', type=lambda x: x.strip())
     parser.add_argument('-p', '--password', dest='password', required=False, default=None,
                         help='The NESSUS password', type=lambda x: x.strip())
-    parser.add_argument('-F', '--fancy', action='store_true', help='Enable colourful logging output')
+    parser.add_argument('-F', '--fancy', action='store_true',
-    parser.add_argument('-d', '--debug', action='store_true', help='Enable debugging messages')
+                        help='Enable colourful logging output')
-    parser.add_argument('--mock', action='store_true', help='Enable mocked API responses')
+    parser.add_argument('-d', '--debug', action='store_true',
-    parser.add_argument('--mock_dir', dest='mock_dir', required=False, default='test',
+                        help='Enable debugging messages')
+    parser.add_argument('--mock', action='store_true',
+                        help='Enable mocked API responses')
+    parser.add_argument('--mock_dir', dest='mock_dir', required=False, default=None,
                         help='Path of test directory')
     args = parser.parse_args()
 

vulnwhisp/frameworks/nessus.py

@@ -1,11 +1,13 @@
-from datetime import datetime
-import sys
-import time
 import json
 import logging
+import sys
+import time
+from datetime import datetime
+
 import pytz
 import requests
 from requests.packages.urllib3.exceptions import InsecureRequestWarning
+
 requests.packages.urllib3.disable_warnings(InsecureRequestWarning)
 
 
@@ -34,7 +36,10 @@ class NessusAPI(object):
         self.base = 'https://{hostname}:{port}'.format(hostname=hostname, port=port)
         self.verbose = verbose
 
-        self.headers = {
+        self.session = requests.Session()
+        self.session.verify = False
+        self.session.stream = True
+        self.session.headers = {
             'Origin': self.base,
             'Accept-Encoding': 'gzip, deflate, br',
             'Accept-Language': 'en-US,en;q=0.8',
@@ -52,27 +57,24 @@ class NessusAPI(object):
         self.scan_ids = self.get_scan_ids()
 
     def login(self):
-        resp = self.get_token()
+        auth = '{"username":"%s", "password":"%s"}' % (self.user, self.password)
+        resp = self.request(self.SESSION, data=auth, json_output=False)
         if resp.status_code == 200:
-            self.headers['X-Cookie'] = 'token={token}'.format(token=resp.json()['token'])
+            self.session.headers['X-Cookie'] = 'token={token}'.format(token=resp.json()['token'])
         else:
             raise Exception('[FAIL] Could not login to Nessus')
 
-    def request(self, url, data=None, headers=None, method='POST', download=False, json=False):
+    def request(self, url, data=None, headers=None, method='POST', download=False, json_output=False):
-        if headers is None:
-            headers = self.headers
         timeout = 0
         success = False
         
+        method = method.lower()
         url = self.base + url
         self.logger.debug('Requesting to url {}'.format(url))
-        methods = {'GET': requests.get,
-                   'POST': requests.post,
-                   'DELETE': requests.delete}
 
         while (timeout <= 10) and (not success):
-            data = methods[method](url, data=data, headers=self.headers, verify=False)
+            response = getattr(self.session, method)(url, data=data)
-            if data.status_code == 401:
+            if response.status_code == 401:
                 if url == self.base + self.SESSION:
                     break
                 try:
@@ -84,20 +86,22 @@ class NessusAPI(object):
             else:
                 success = True
 
-        if json:
+        if json_output:
-            data = data.json()
+            return response.json()
         if download:
             self.logger.debug('Returning data.content')
-            return data.content
+            response_data = ''
-        return data
+            count = 0
-
+            for chunk in response.iter_content(chunk_size=8192):
-    def get_token(self):
+                count += 1
-        auth = '{"username":"%s", "password":"%s"}' % (self.user, self.password)
+                if chunk:
-        token = self.request(self.SESSION, data=auth, json=False)
+                    response_data += chunk
-        return token
+            self.logger.debug('Processed {} chunks'.format(count))
+            return response_data
+        return response
 
     def get_scans(self):
-        scans = self.request(self.SCANS, method='GET', json=True)
+        scans = self.request(self.SCANS, method='GET', json_output=True)
         return scans
 
     def get_scan_ids(self):
@@ -107,10 +111,10 @@ class NessusAPI(object):
         return scan_ids
 
     def get_scan_history(self, scan_id):
-        data = self.request(self.SCAN_ID.format(scan_id=scan_id), method='GET', json=True)
+        data = self.request(self.SCAN_ID.format(scan_id=scan_id), method='GET', json_output=True)
         return data['history']
 
-    def download_scan(self, scan_id=None, history=None, export_format="", chapters="", dbpasswd="", profile=""):
+    def download_scan(self, scan_id=None, history=None, export_format="", profile=""):
         running = True
         counter = 0
 
@@ -120,7 +124,7 @@ class NessusAPI(object):
         else:
             query = self.EXPORT_HISTORY.format(scan_id=scan_id, history_id=history)
             scan_id = str(scan_id)
-        req = self.request(query, data=json.dumps(data), method='POST', json=True)
+        req = self.request(query, data=json.dumps(data), method='POST', json_output=True)
         try:
             file_id = req['file']
             token_id = req['token'] if 'token' in req else req['temp_token']
@@ -131,7 +135,7 @@ class NessusAPI(object):
             time.sleep(2)
             counter += 2
             report_status = self.request(self.EXPORT_STATUS.format(scan_id=scan_id, file_id=file_id), method='GET',
-                                         json=True)
+                                         json_output=True)
             running = report_status['status'] != 'ready'
             sys.stdout.write(".")
             sys.stdout.flush()

vulnwhisp/frameworks/qualys_vuln.py

@@ -2,12 +2,13 @@
 # -*- coding: utf-8 -*-
 __author__ = 'Nathan Young'
 
-import xml.etree.ElementTree as ET
-import sys
 import logging
-import qualysapi
+import sys
-import pandas as pd
+import xml.etree.ElementTree as ET
+
 import dateutil.parser as dp
+import pandas as pd
+import qualysapi
 
 
 class qualysWhisperAPI(object):
@@ -109,7 +110,10 @@ class qualysVulnScan:
         self.logger.info('Downloading scan ID: {}'.format(scan_id))
         scan_report = self.qw.get_scan_details(scan_id=scan_id)
         if not scan_report.empty:
-            keep_columns = ['category', 'cve_id', 'cvss3_base', 'cvss3_temporal', 'cvss_base', 'cvss_temporal', 'dns', 'exploitability', 'fqdn', 'impact', 'ip', 'ip_status', 'netbios', 'os', 'pci_vuln', 'port', 'protocol', 'qid', 'results', 'severity', 'solution', 'ssl', 'threat', 'title', 'type', 'vendor_reference']
+            keep_columns = ['category', 'cve_id', 'cvss3_base', 'cvss3_temporal', 'cvss_base',
+                'cvss_temporal', 'dns', 'exploitability', 'fqdn', 'impact', 'ip', 'ip_status',
+                'netbios', 'os', 'pci_vuln', 'port', 'protocol', 'qid', 'results', 'severity',
+                'solution', 'ssl', 'threat', 'title', 'type', 'vendor_reference']
             scan_report = scan_report.filter(keep_columns)
             scan_report['severity'] = scan_report['severity'].astype(int).astype(str)
             scan_report['qid'] = scan_report['qid'].astype(int).astype(str)

vulnwhisp/test/mock.py

@@ -6,12 +6,12 @@ import httpretty
 class mockAPI(object):
     def __init__(self, mock_dir=None, debug=False):
         self.mock_dir = mock_dir
+
         if not self.mock_dir:
             # Try to guess the mock_dir if python setup.py develop was used
-            self.mock_dir = '/'.join(__file__.split('/')[:-3]) + '/test'
+            self.mock_dir = '/'.join(__file__.split('/')[:-3]) + '/tests/data'
 
         self.logger = logging.getLogger('mockAPI')
-
         if debug:
             self.logger.setLevel(logging.DEBUG)
 

vulnwhisp/vulnwhisp.py

@@ -410,7 +410,7 @@ class vulnWhispererNessus(vulnWhispererBase):
                 if status in ['completed', 'imported']:
                     file_name = '%s_%s_%s_%s.%s' % (scan_name, scan_id,
                                                     history_id, norm_time, 'csv')
-                    repls = (('\\', '_'), ('/', '_'), ('/', '_'), (' ', '_'))
+                    repls = (('\\', '_'), ('/', '_'), (' ', '_'))
                     file_name = reduce(lambda a, kv: a.replace(*kv), repls, file_name)
                     relative_path_name = self.path_check(folder_name + '/' + file_name)