Small style changes for PR 75

pyt/__main__.py

@@ -13,10 +13,11 @@
 from .fixed_point import analyse
 from .framework_adaptor import FrameworkAdaptor
 from .framework_helper import (
+    is_django_view_function,
     is_flask_route_function,
     is_function,
-    is_function_without_leading_,
-    is_django_view_function)
+    is_function_without_leading_
+)
 from .github_search import scan_github, set_github_api_token
 from .interprocedural_cfg import interprocedural
 from .intraprocedural_cfg import intraprocedural
@@ -83,7 +84,7 @@ def parse_args(args):
                         help='Choose logging level: CRITICAL, ERROR,' +
                         ' WARNING(Default), INFO, DEBUG, NOTSET.', type=str)
     parser.add_argument('-a', '--adaptor',
-                        help='Choose an adaptor: Flask(Default) or Every or Pylons or Django.',
+                        help='Choose an adaptor: Flask(Default), Django, Every or Pylons',
                         type=str)
     parser.add_argument('-db', '--create-database',
                         help='Creates a sql file that can be used to' +
@@ -222,14 +223,14 @@ def main(command_line_args=sys.argv[1:]):
                                               local_modules,
                                               path)
         cfg_list.append(interprocedural_cfg)
-        if args.adaptor and args.adaptor.lower().startswith('e'):
-            framework_route_criteria = is_function
-        elif args.adaptor and args.adaptor.lower().startswith('p'):
-            framework_route_criteria = is_function_without_leading_
-        elif args.adaptor and args.adaptor.lower().startswith('d'):
-            framework_route_criteria = is_django_view_function
-        else:
-            framework_route_criteria = is_flask_route_function
+        framework_route_criteria = is_flask_route_function
+        if args.adaptor:
+            if args.adaptor.lower().startswith('e'):
+                framework_route_criteria = is_function
+            elif args.adaptor.lower().startswith('p'):
+                framework_route_criteria = is_function_without_leading_
+            elif args.adaptor.lower().startswith('d'):
+                framework_route_criteria = is_django_view_function
         # Add all the route functions to the cfg_list
         FrameworkAdaptor(cfg_list, project_modules, local_modules, framework_route_criteria)
 

pyt/framework_helper.py

@@ -1,8 +1,7 @@
 """Provides helper functions that help with determining if a function is a route function."""
 import ast
 
-from pyt.base_cfg import Function
-from .ast_helper import get_call_names, Arguments
+from .ast_helper import get_call_names
 
 
 def is_function(function):

tests/framework_helper_test.py

@@ -1,15 +1,16 @@
 from .base_test_case import BaseTestCase
 from pyt.framework_adaptor import _get_func_nodes
 from pyt.framework_helper import (
+    is_django_view_function,
     is_flask_route_function,
     is_function,
     is_function_without_leading_,
-    is_django_view_function
 )
 
+
 class FrameworkEngineTest(BaseTestCase):
     def test_find_flask_functions(self):
-        self.cfg_create_from_file('example/example_inputs/flask_function_and_normal_functions.py')
+        self.cfg_create_from_file('example/example_inputs/django_flask_and_normal_functions.py')
 
         cfg_list = [self.cfg]
         funcs = _get_func_nodes()
@@ -24,7 +25,7 @@ def test_find_flask_functions(self):
 
 
     def test_find_every_function_without_leading_underscore(self):
-        self.cfg_create_from_file('example/example_inputs/flask_function_and_normal_functions.py')
+        self.cfg_create_from_file('example/example_inputs/django_flask_and_normal_functions.py')
 
         cfg_list = [self.cfg]
         funcs = _get_func_nodes()
@@ -37,7 +38,7 @@ def test_find_every_function_without_leading_underscore(self):
         self.assertEqual(i, 3)
 
     def test_find_every_function(self):
-        self.cfg_create_from_file('example/example_inputs/flask_function_and_normal_functions.py')
+        self.cfg_create_from_file('example/example_inputs/django_flask_and_normal_functions.py')
 
         cfg_list = [self.cfg]
         funcs = _get_func_nodes()
@@ -50,7 +51,7 @@ def test_find_every_function(self):
         self.assertEqual(len(funcs), 4)
 
     def test_find_django_functions(self):
-        self.cfg_create_from_file('example/example_inputs/flask_function_and_normal_functions.py')
+        self.cfg_create_from_file('example/example_inputs/django_flask_and_normal_functions.py')
 
         cfg_list = [self.cfg]
         funcs = _get_func_nodes()

tests/vulnerabilities_test.py

@@ -6,7 +6,7 @@
 from pyt.constraint_table import constraint_table, initialize_constraint_table
 from pyt.fixed_point import analyse
 from pyt.framework_adaptor import FrameworkAdaptor
-from pyt.framework_helper import is_flask_route_function, is_django_view_function
+from pyt.framework_helper import is_django_view_function, is_flask_route_function
 from pyt.lattice import Lattice
 from pyt.reaching_definitions_taint import ReachingDefinitionsTaintAnalysis
 
@@ -318,7 +318,7 @@ def test_XSS_url_result(self):
         vulnerability_description = str(vulnerability_log.vulnerabilities[0])
         EXPECTED_VULNERABILITY_DESCRIPTION = """
             File: example/vulnerable_code/XSS_url.py
-             > User input at line 4, trigger word "Framework function URL parameter": 
+             > User input at line 4, trigger word "Framework function URL parameter":
                 url
             Reassigned in: 
                 File: example/vulnerable_code/XSS_url.py
@@ -480,13 +480,13 @@ def test_django_view_param(self):
 
         EXPECTED_VULNERABILITY_DESCRIPTION = """
             File: example/vulnerable_code/django_XSS.py
-             > User input at line 4, trigger word "Framework function URL parameter": 
+             > User input at line 4, trigger word "Framework function URL parameter":
                 param
-            Reassigned in: 
+            Reassigned in:
                 File: example/vulnerable_code/django_XSS.py
                  > Line 5: ret_xss1 = ¤call_1
             File: example/vulnerable_code/django_XSS.py
-             > reaches line 5, trigger word "render(": 
+             > reaches line 5, trigger word "render(":
                 ¤call_1 = ret_render(request, 'templates/xss.html', 'param'param)
         """
         self.assertTrue(self.string_compare_alpha(vulnerability_description, EXPECTED_VULNERABILITY_DESCRIPTION))