Merge pull request #109 from Santandersecurityresearch/develop

Changes for v1.2.0

Author: javixeneize

README.md

@@ -1,5 +1,5 @@
 [![GitHub release](https://img.shields.io/github/release/Santandersecurityresearch/DrHeader.svg)](https://GitHub.com/Santandersecurityresearch/DrHeader/releases/)
-[![GitHub commits](https://img.shields.io/github/commits-since/Santandersecurityresearch/DrHeader/v1.1.1.svg)](https://GitHub.com/Santandersecurityresearch/DrHeader/commit/)
+[![GitHub commits](https://img.shields.io/github/commits-since/Santandersecurityresearch/DrHeader/v1.2.0.svg)](https://GitHub.com/Santandersecurityresearch/DrHeader/commit/)
 [![Github all releases](https://img.shields.io/github/downloads/Santandersecurityresearch/DrHeader/total.svg)](https://GitHub.com/Santandersecurityresearch/DrHeader/releases/)
 [![HitCount](http://hits.dwyl.io/Santandersecurityresearch/DrHeader.svg)](http://hits.dwyl.io/Santandersecurityresearch/DrHeader)
 [![Total alerts](https://img.shields.io/lgtm/alerts/g/Santandersecurityresearch/DrHeader.svg?logo=lgtm&logoWidth=18)](https://lgtm.com/projects/g/Santandersecurityresearch/DrHeader/alerts/)
@@ -68,6 +68,7 @@ There are a number of parameters you can specify during bulk scans, these are:
 | --json            | Output report as json                                  |
 | --debug           | Show error messages                                    |
 | --rules FILENAME  | Use custom rule set                                    |
+| --rules-uri URL   | Use custom rule set, to download from a remote server  |
 | --merge           | Merge custom rule set on top of default set            |
 | --help            | Show this message and exit                             |
 | --junit           | Creates a junit report in `./reports/junit.xml` folder |

drheader/cli.py

@@ -8,12 +8,17 @@
 import click
 import jsonschema
 import sys
+import os
 import validators
 from tabulate import tabulate
 
 from drheader import Drheader
 from drheader.cli_utils import echo_bulk_report, file_junit_report
-from drheader.utils import load_rules
+from drheader.utils import load_rules, get_rules_from_uri
+
+
+EXIT_CODE_NO_ERROR = os.EX_OK
+EXIT_CODE_FAILURE = os.EX_SOFTWARE
 
 
 @click.group()
@@ -31,8 +36,9 @@ def scan():
 @click.option('--json', 'json_output', help='Output report as json', is_flag=True)
 @click.option('--debug', help='Show error messages', is_flag=True)
 @click.option('--rules', 'rule_file', help='Use custom rule set', type=click.File())
+@click.option('--rules-uri', 'rule_uri', help='Use custom rule set, downloaded from URI')
 @click.option('--merge', help='Merge custom file rules, on top of default rules', is_flag=True)
-def compare(file, json_output, debug, rule_file, merge):
+def compare(file, json_output, debug, rule_file, rule_uri, merge):
     """
     If you have headers you would like to test with drheader, you can "compare" them with your ruleset this command.
 
@@ -60,7 +66,7 @@ def compare(file, json_output, debug, rule_file, merge):
             ...
         ]
     """
-
+    exit_code = EXIT_CODE_NO_ERROR
     audit = []
     schema = {
         "type": "array",
@@ -90,38 +96,64 @@ def compare(file, json_output, debug, rule_file, merge):
     except Exception as e:
         raise click.ClickException(e)
 
+    if rule_uri and not rule_file:
+        if not validators.url(rule_uri):
+            raise click.ClickException(message='"{}" is not a valid URL.'.format(rule_uri))
+        try:
+            rule_file = get_rules_from_uri(rule_uri)
+        except Exception as e:
+            if debug:
+                raise click.ClickException(e)
+            else:
+                raise click.ClickException('No content retrieved from rules-uri.')
+
     rules = load_rules(rule_file, merge)
 
     for i in data:
         logging.debug('Analysing : {}'.format(i['url']))
         drheader_instance = Drheader(url=i['url'], headers=i['headers'])
         drheader_instance.analyze(rules)
         audit.append({'url': i['url'], 'report': drheader_instance.report})
+        if drheader_instance.report:
+            exit_code = EXIT_CODE_FAILURE
 
     echo_bulk_report(audit, json_output)
+    sys.exit(exit_code)
 
 
 @scan.command()
 @click.argument('target_url', required=True)
 @click.option('--json', 'json_output', help='Output report as json', is_flag=True)
 @click.option('--debug', help='Show error messages', is_flag=True)
 @click.option('--rules', 'rule_file', help='Use custom rule set', type=click.File())
+@click.option('--rules-uri', 'rule_uri', help='Use custom rule set, downloaded from URI')
 @click.option('--merge', help='Merge custom file rules, on top of default rules', is_flag=True)
 @click.option('--junit', help='Produces a junit report with the result of the scan', is_flag=True)
-def single(target_url, json_output, debug, rule_file, merge, junit):
+def single(target_url, json_output, debug, rule_file, rule_uri, merge, junit):
     """
     Scan a single http(s) endpoint with drheader.
 
     NOTE: URL parameters are currently only supported on bulk scans.
     """
-
+    exit_code = EXIT_CODE_NO_ERROR
     if debug:
         logging.basicConfig(level=logging.DEBUG)
 
     logging.debug('Validating: {}'.format(target_url))
     if not validators.url(target_url):
         raise click.ClickException(message='"{}" is not a valid URL.'.format(target_url))
 
+    if rule_uri and not rule_file:
+        if not validators.url(rule_uri):
+            raise click.ClickException(message='"{}" is not a valid URL.'.format(rule_uri))
+        try:
+            rule_file = get_rules_from_uri(rule_uri)
+        except Exception as e:
+            if debug:
+                raise click.ClickException(e)
+            else:
+                raise click.ClickException('No content retrieved from rules-uri.')
+
     rules = load_rules(rule_file, merge)
 
     try:
@@ -142,6 +174,9 @@ def single(target_url, json_output, debug, rule_file, merge, junit):
         else:
             raise click.ClickException('Failed to analyze headers.')
 
+    if drheader_instance.report:
+        exit_code = EXIT_CODE_FAILURE
+
     if json_output:
         click.echo(json.dumps(drheader_instance.report))
     else:
@@ -158,7 +193,7 @@ def single(target_url, json_output, debug, rule_file, merge, junit):
                 click.echo(tabulate(values, tablefmt="presto"))
     if junit:
         file_junit_report(rules, drheader_instance.report)
-    return 0
+    sys.exit(exit_code)
 
 
 @scan.command()
@@ -169,8 +204,9 @@ def single(target_url, json_output, debug, rule_file, merge, junit):
 @click.option('--json', 'json_output', help='Output report as json', is_flag=True)
 @click.option('--debug', help='Show error messages', is_flag=True)
 @click.option('--rules', 'rule_file', help='Use custom rule set', type=click.File())
+@click.option('--rules-uri', 'rule_uri', help='Use custom rule set, downloaded from URI')
 @click.option('--merge', help='Merge custom file rules, on top of default rules', is_flag=True)
-def bulk(file, json_output, post, input_format, debug, rule_file, merge):
+def bulk(file, json_output, post, input_format, debug, rule_file, rule_uri, merge):
     """
     Scan multiple http(s) endpoints with drheader.
 
@@ -195,7 +231,7 @@ def bulk(file, json_output, post, input_format, debug, rule_file, merge):
 
     NOTE: URL parameters are currently only supported on bulk scans.
     """
-
+    exit_code = EXIT_CODE_NO_ERROR
     audit = []
     urls = []
     schema = {
@@ -235,6 +271,17 @@ def bulk(file, json_output, post, input_format, debug, rule_file, merge):
 
     logging.debug('Found {} URLs'.format(len(urls)))
 
+    if rule_uri and not rule_file:
+        if not validators.url(rule_uri):
+            raise click.ClickException(message='"{}" is not a valid URL.'.format(rule_uri))
+        try:
+            rule_file = get_rules_from_uri(rule_uri)
+        except Exception as e:
+            if debug:
+                raise click.ClickException(e)
+            else:
+                raise click.ClickException('No content retrieved from rules-uri.')
+
     rules = load_rules(rule_file, merge)
 
     for i, v in enumerate(urls):
@@ -243,9 +290,11 @@ def bulk(file, json_output, post, input_format, debug, rule_file, merge):
         logging.debug('Analysing: {}...'.format(v))
         drheader_instance.analyze(rules)
         audit.append({'url': v['url'], 'report': drheader_instance.report})
+        if drheader_instance.report:
+            exit_code = EXIT_CODE_FAILURE
 
     echo_bulk_report(audit, json_output)
-    return 0
+    sys.exit(exit_code)
 
 
 if __name__ == "__main__":

drheader/utils.py

@@ -4,7 +4,8 @@
 
 import logging
 import os
-
+import io
+import requests
 import yaml
 
 
@@ -61,3 +62,18 @@ def merge_rules(default_rules, custom_rules):
         default_rules['Headers'][rule] = custom_rules['Headers'][rule]
 
     return default_rules
+
+
+def get_rules_from_uri(URI):
+    """
+    Retrieves custom rule set from URL
+    :param URI: URL to your custom rules file
+    :type URI: uri
+    :return: rules file
+    :rtype: file
+    """
+    download = requests.get(URI)
+    if not download.content:
+        raise Exception('No content retrieved from {}'.format(URI))
+    file = io.BytesIO(download.content)
+    return file

requirements_dev.txt

@@ -16,3 +16,4 @@ validators==0.14.0
 unittest2==1.1.0
 xmlunittest==0.5.0
 junit-xml==1.9
+responses==0.10.12

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 1.1.1
+current_version = 1.2.0
 commit = True
 tag = True
 

setup.py

@@ -49,6 +49,6 @@
     test_suite='tests',
     tests_require=test_requirements, 
     url='https://github.com/santandersecurityresearch/drheader',
-    version='1.1.1',
+    version='1.2.0',
     zip_safe=False,
 )

tests/test_utils.py

@@ -6,8 +6,9 @@
 import os
 import yaml
 import unittest2
+import responses
 
-from drheader.utils import load_rules
+from drheader.utils import load_rules, get_rules_from_uri
 
 
 class TestUtilsFunctions(unittest2.TestCase):
@@ -45,6 +46,19 @@ def test_load_rules_bad_parameter(self):
         with self.assertRaises(AttributeError):
             load_rules(2)
 
+    @responses.activate
+    def test_get_rules_from_uri_wrong_URI(self):
+        responses.add(responses.GET, 'http://mydomain.com/custom.yml', status=404)
+        with self.assertRaises(Exception):
+            get_rules_from_uri("http://mydomain.com/custom.yml")
+
+    @responses.activate
+    def test_get_rules_from_uri_good_URI(self):
+        responses.add(responses.GET, 'http://localhost:8080/custom.yml', json=self.custom_rules, status=200)
+        file = get_rules_from_uri("http://localhost:8080/custom.yml")
+        content = yaml.safe_load(file.read())
+        self.assertEqual(content, self.custom_rules)
+
 
 # start unittest2 to run these tests
 if __name__ == "__main__":