KevinHock · Jul 14, 2017 · Jun 8, 2017 · Jun 8, 2017 · Jun 8, 2017 · Jun 9, 2017
diff --git a/...uts/flask_function_and_normal_function.py → ...ts/flask_function_and_normal_functions.py b/...uts/flask_function_and_normal_function.py → ...ts/flask_function_and_normal_functions.py
@@ -1,6 +1,9 @@
 def foo():
     print('h')
 
+def _hidden_foo():
+    print('h')
+
 @app.route('/', methods = ['GET'])
 def flask_function(x):
     return x

diff --git a/example/vulnerable_code/tainted_arg_normal_function.py b/example/vulnerable_code/tainted_arg_normal_function.py
@@ -0,0 +1,20 @@
+import os
+
+
+def store_uploaded_file(title, uploaded_file):
+    """ Stores a temporary uploaded file on disk """
+    upload_dir_path = '%s/static/taskManager/uploads' % (
+        os.path.dirname(os.path.realpath(__file__)))
+    if not os.path.exists(upload_dir_path):
+        os.makedirs(upload_dir_path)
+
+    # A1: Injection (shell)
+    os.system(
+        "mv " +
+        uploaded_file.temporary_file_path() +
+        " " +
+        "%s/%s" %
+        (upload_dir_path,
+         title))
+
+    return '/static/taskManager/uploads/%s' % (title)
diff --git a/pyt/__main__.py b/pyt/__main__.py
@@ -11,7 +11,12 @@
 from .draw import draw_cfgs, draw_lattices
 from .constraint_table import initialize_constraint_table, print_table
 from .fixed_point import analyse
-from .flask_adaptor import FlaskAdaptor
+from .framework_adaptor import FrameworkAdaptor
+from .framework_helper import (
+    is_flask_route_function,
+    is_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
@@ -78,7 +83,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 DJANGO.',
+                        help='Choose an adaptor: Flask(Default) or Every or Pylons.',
                         type=str)
     parser.add_argument('-db', '--create-database',
                         help='Creates a sql file that can be used to' +
@@ -212,11 +217,19 @@ def main(command_line_args=sys.argv[1:]):
     if args.intraprocedural_analysis:
         intraprocedural(project_modules, cfg_list)
     else:
-        cfg_list.append(interprocedural(tree,
-                                        project_modules,
-                                        local_modules,
-                                        path))
-        adaptor_type = FlaskAdaptor(cfg_list, project_modules, local_modules)
+        interprocedural_cfg = interprocedural(tree,
+                                              project_modules,
+                                              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_
+        else:
+            framework_route_criteria = is_flask_route_function
+        # Add all the route functions to the cfg_list
+        FrameworkAdaptor(cfg_list, project_modules, local_modules, framework_route_criteria)
 
     initialize_constraint_table(cfg_list)
 

diff --git a/pyt/flask_adaptor.py b/pyt/flask_adaptor.py
diff --git a/pyt/framework_adaptor.py b/pyt/framework_adaptor.py
@@ -1,23 +1,77 @@
-"""A framework adaptor is a adaptor used to adapt
-the source code to a specific framework."""
-from abc import ABCMeta, abstractmethod
+"""A generic framework adaptor that leaves route criteteria to the caller."""
+import ast
 
+from .ast_helper import Arguments
 from .base_cfg import AssignmentNode
+from .interprocedural_cfg import interprocedural
+from .module_definitions import project_definitions
 
 
-class FrameworkAdaptor(metaclass=ABCMeta):
-    """An engine that should be used as base class
-    to specify how to find all sources and sinks."""
+class FrameworkAdaptor():
+    """An engine that uses the template pattern to find all
+    entry points in a framework and then taints their arguments.
+    """
 
-    def __init__(self, cfg_list, project_modules, local_modules):
+    def __init__(self, cfg_list, project_modules, local_modules, is_route_function):
         self.cfg_list = cfg_list
         self.project_modules = project_modules
         self.local_modules = local_modules
+        self.is_route_function = is_route_function
         self.run()
 
-    @abstractmethod
+    def get_func_cfg_with_tainted_args(self, definition):
+        """Build a function cfg and return it, with all arguments tainted."""
+        func_cfg = interprocedural(definition.node, self.project_modules,
+                                   self.local_modules, definition.path,
+                                   definition.module_definitions)
+
+        args = Arguments(definition.node.args)
+        if args:
+            function_entry_node = func_cfg.nodes[0]
+            function_entry_node.outgoing = []
+            first_node_after_args = func_cfg.nodes[1]
+            first_node_after_args.ingoing = []
+
+            # We're just gonna give all the tainted args the lineno of the def
+            definition_lineno = definition.node.lineno
+
+            # Taint all the arguments
+            for arg in args:
+                tainted_node = TaintedNode(arg, arg,
+                                           None, [],
+                                           line_number=definition_lineno,
+                                           path=definition.path)
+                function_entry_node.connect(tainted_node)
+                # 1 and not 0 so that Entry Node remains first in the list
+                func_cfg.nodes.insert(1, tainted_node)
+
+            first_arg = func_cfg.nodes[len(args)]
+            first_arg.connect(first_node_after_args)
+
+        return func_cfg
+
+    def find_route_functions_taint_args(self):
+        """Find all route functions and taint all of their arguments.
+
+        Yields:
+            CFG of each route function, with args marked as tainted.
+        """
+        for definition in _get_func_nodes():
+            if self.is_route_function(definition.node):
+                yield self.get_func_cfg_with_tainted_args(definition)
+
     def run(self):
-        pass
+        """Run find_route_functions_taint_args on each CFG."""
+        function_cfgs = list()
+        for _ in self.cfg_list:
+            function_cfgs.extend(self.find_route_functions_taint_args())
+        self.cfg_list.extend(function_cfgs)
+
+
+def _get_func_nodes():
+    """Get all function nodes."""
+    return [definition for definition in project_definitions.values()
+            if isinstance(definition.node, ast.FunctionDef)]
 
 
 class TaintedNode(AssignmentNode):

diff --git a/pyt/framework_helper.py b/pyt/framework_helper.py
@@ -0,0 +1,32 @@
+"""Provides helper functions that help with determining if a function is a route function."""
+import ast
+
+from .ast_helper import get_call_names
+
+
+def is_function(function):
+    """Always returns true because arg is always a function."""
+    return True
+
+
+def is_flask_route_function(ast_node):
+    """Check whether function uses a route decorator."""
+    for decorator in ast_node.decorator_list:
+        if isinstance(decorator, ast.Call):
+            if _get_last_of_iterable(get_call_names(decorator.func)) == 'route':
+                return True
+    return False
+
+
+def is_function_without_leading_(ast_node):
+    if ast_node.name.startswith('_'):
+        return False
+    return True
+
+
+def _get_last_of_iterable(iterable):
+    """Get last element of iterable."""
+    item = None
+    for item in iterable:
+        pass
+    return item
diff --git a/pyt/trigger_definitions/flask_trigger_words.pyt b/pyt/trigger_definitions/flask_trigger_words.pyt
@@ -12,6 +12,7 @@ sinks:
 replace( -> escape
 send_file( -> '..', '..' in
 execute(
+system(
 filter(
 subprocess.call(
 render_template(

diff --git a/tests/flask_adaptor_test.py b/tests/flask_adaptor_test.py
diff --git a/tests/framework_helper_test.py b/tests/framework_helper_test.py
@@ -0,0 +1,49 @@
+from .base_test_case import BaseTestCase
+from pyt.framework_adaptor import _get_func_nodes
+from pyt.framework_helper import (
+    is_flask_route_function,
+    is_function,
+    is_function_without_leading_
+)
+
+class FrameworkEngineTest(BaseTestCase):
+    def test_find_flask_functions(self):
+        self.cfg_create_from_file('example/example_inputs/flask_function_and_normal_functions.py')
+
+        cfg_list = [self.cfg]
+        funcs = _get_func_nodes()
+
+        i = 0
+        for func in funcs:
+            if is_flask_route_function(func.node):
+                self.assertEqual(func.node.name, 'flask_function')
+                i = i + 1
+        # So it is supposed to be 1, because foo is not an app.route
+        self.assertEqual(i, 1)
+
+
+    def test_find_every_function_without_leading_underscore(self):
+        self.cfg_create_from_file('example/example_inputs/flask_function_and_normal_functions.py')
+
+        cfg_list = [self.cfg]
+        funcs = _get_func_nodes()
+
+        i = 0
+        for func in funcs:
+            if is_function_without_leading_(func.node):
+                i = i + 1
+        # So it is supposed to be 2, because we count all functions without a leading underscore
+        self.assertEqual(i, 2)
+
+    def test_find_every_function(self):
+        self.cfg_create_from_file('example/example_inputs/flask_function_and_normal_functions.py')
+
+        cfg_list = [self.cfg]
+        funcs = _get_func_nodes()
+
+        i = 0
+        for func in funcs:
+            if is_function(func.node):
+                i = i + 1
+        # So it is supposed to be 3, because we count all functions
+        self.assertEqual(len(funcs), 3)
diff --git a/tests/vulnerabilities_across_files_test.py b/tests/vulnerabilities_across_files_test.py
@@ -7,7 +7,8 @@
 from pyt.base_cfg import Node
 from pyt.constraint_table import constraint_table, initialize_constraint_table
 from pyt.fixed_point import analyse
-from pyt.flask_adaptor import FlaskAdaptor
+from pyt.framework_adaptor import FrameworkAdaptor
+from pyt.framework_helper import is_flask_route_function
 from pyt.lattice import Lattice
 from pyt.project_handler import get_directory_modules, get_modules
 from pyt.reaching_definitions_taint import ReachingDefinitionsTaintAnalysis
@@ -24,7 +25,7 @@ def run_analysis(self, path):
 
         cfg_list = [self.cfg]
 
-        FlaskAdaptor(cfg_list, [], [])
+        FrameworkAdaptor(cfg_list, [], [], is_flask_route_function)
 
         initialize_constraint_table(cfg_list)