Use import_alias_mapping for blackbox and built-in aliases, as well

* This will give fully qualified names for blackboxes like flask
* Improve readability by using keyword arguments

Author: wchresta

examples/vulnerable_code/command_injection_with_aliases.py

@@ -4,9 +4,28 @@
 from os import system as mysystem
 from subprocess import call as mycall, Popen as mypopen
 
-os.system("ls")
-myos.system("ls")
-system("ls")
-mysystem("ls")
-mycall("ls")
-mypopen("ls")
+from flask import Flask, render_template, request
+
+app = Flask(__name__)
+
+
+@app.route('/menu', methods=['POST'])
+def menu():
+    param = request.form['suggestion']
+    command = 'echo ' + param + ' >> ' + 'menu.txt'
+
+    os.system(command)
+    myos.system(command)
+    system(command)
+    mysystem(command)
+    mycall(command)
+    mypopen(command)
+
+    with open('menu.txt', 'r') as f:
+        menu_ctx = f.read()
+
+    return render_template('command_injection.html', menu=menu_ctx)
+
+
+if __name__ == '__main__':
+    app.run(debug=True)

pyt/cfg/stmt_visitor.py

@@ -62,7 +62,6 @@
 class StmtVisitor(ast.NodeVisitor):
     def __init__(self, allow_local_directory_imports=True):
         self._allow_local_modules = allow_local_directory_imports
-        self.bb_or_bi_aliases = {}
         super().__init__()
 
     def visit_Module(self, node):
@@ -628,7 +627,8 @@ def add_blackbox_or_builtin_call(self, node, blackbox):  # noqa: C901
 
         # Check if function call matches a blackbox/built-in alias and if so, resolve it
         # This resolves aliases like "from os import system as mysys" as: mysys -> os.system
-        call_function_label = fully_qualify_alias_labels(call_function_label, self.bb_or_bi_aliases)
+        local_definitions = self.module_definitions_stack[-1]
+        call_function_label = fully_qualify_alias_labels(call_function_label, local_definitions.import_alias_mapping)
 
         # Create e.g. ~call_1 = ret_func_foo
         LHS = CALL_IDENTIFIER + 'call_' + str(saved_function_call_index)
@@ -924,10 +924,10 @@ def from_directory_import(
         if init_exists and not skip_init:
             package_name = os.path.split(module_path)[1]
             return self.add_module(
-                (module[0], init_file_location),
-                package_name,
-                local_names,
-                import_alias_mapping,
+                module=(module[0], init_file_location),
+                module_or_package_name=package_name,
+                local_names=local_names,
+                import_alias_mapping=import_alias_mapping,
                 is_init=True,
                 from_from=True
             )
@@ -937,10 +937,10 @@ def from_directory_import(
                 new_init_file_location = os.path.join(full_name, '__init__.py')
                 if os.path.isfile(new_init_file_location):
                     self.add_module(
-                        (real_name, new_init_file_location),
-                        real_name,
-                        local_names,
-                        import_alias_mapping,
+                        module=(real_name, new_init_file_location),
+                        module_or_package_name=real_name,
+                        local_names=local_names,
+                        import_alias_mapping=import_alias_mapping,
                         is_init=True,
                         from_from=True,
                         from_fdid=True
@@ -950,10 +950,10 @@ def from_directory_import(
             else:
                 file_module = (real_name, full_name + '.py')
                 self.add_module(
-                    file_module,
-                    real_name,
-                    local_names,
-                    import_alias_mapping,
+                    module=file_module,
+                    module_or_package_name=real_name,
+                    local_names=local_names,
+                    import_alias_mapping=import_alias_mapping,
                     from_from=True
                 )
         return IgnoredNode()
@@ -964,10 +964,10 @@ def import_package(self, module, module_name, local_name, import_alias_mapping):
         init_exists = os.path.isfile(init_file_location)
         if init_exists:
             return self.add_module(
-                (module[0], init_file_location),
-                module_name,
-                local_name,
-                import_alias_mapping,
+                module=(module[0], init_file_location),
+                module_or_package_name=module_name,
+                local_names=local_name,
+                import_alias_mapping=import_alias_mapping,
                 is_init=True
             )
         else:
@@ -1010,10 +1010,10 @@ def handle_relative_import(self, node):
         # Is it a file?
         if name_with_dir.endswith('.py'):
             return self.add_module(
-                (node.module, name_with_dir),
-                None,
-                as_alias_handler(node.names),
-                retrieve_import_alias_mapping(node.names),
+                module=(node.module, name_with_dir),
+                module_or_package_name=None,
+                local_names=as_alias_handler(node.names),
+                import_alias_mapping=retrieve_import_alias_mapping(node.names),
                 from_from=True
             )
         return self.from_directory_import(
@@ -1036,10 +1036,10 @@ def visit_Import(self, node):
                             retrieve_import_alias_mapping(node.names)
                         )
                     return self.add_module(
-                        module,
-                        name.name,
-                        name.asname,
-                        retrieve_import_alias_mapping(node.names)
+                        module=module,
+                        module_or_package_name=name.name,
+                        local_names=name.asname,
+                        import_alias_mapping=retrieve_import_alias_mapping(node.names)
                     )
             for module in self.project_modules:
                 if name.name == module[0]:
@@ -1051,68 +1051,69 @@ def visit_Import(self, node):
                             retrieve_import_alias_mapping(node.names)
                         )
                     return self.add_module(
-                        module,
-                        name.name,
-                        name.asname,
-                        retrieve_import_alias_mapping(node.names)
+                        module=module,
+                        module_or_package_name=name.name,
+                        local_names=name.asname,
+                        import_alias_mapping=retrieve_import_alias_mapping(node.names)
                     )
         for alias in node.names:
-            if alias.name in uninspectable_modules:
-                # The module is uninspectable (so blackbox or built-in). If it has an alias, we remember
-                # the alias so we can do fully qualified name resolution for blackbox- and built-in trigger words
-                # e.g. we want a call to "os.system" be recognised, even if we do "import os as myos"
-                if alias.asname is not None and alias.asname != alias.name:
-                    self.bb_or_bi_aliases[alias.asname] = alias.name
-            else:
-                log.warn("Cannot inspect module %s", alias.name)
+            # The module is uninspectable (so blackbox or built-in). If it has an alias, we remember
+            # the alias so we can do fully qualified name resolution for blackbox- and built-in trigger words
+            # e.g. we want a call to "os.system" be recognised, even if we do "import os as myos"
+            if alias.asname is not None and alias.asname != alias.name:
+                local_definitions = self.module_definitions_stack[-1]
+                local_definitions.import_alias_mapping[name.asname] = alias.name
+            if alias.name not in uninspectable_modules:
+                log.warning("Cannot inspect module %s", alias.name)
                 uninspectable_modules.add(alias.name)  # Don't repeatedly warn about this
         return IgnoredNode()
 
     def visit_ImportFrom(self, node):
         # Is it relative?
         if node.level > 0:
             return self.handle_relative_import(node)
-        else:
-            for module in self.local_modules:
-                if node.module == module[0]:
-                    if os.path.isdir(module[1]):
-                        return self.from_directory_import(
-                            module,
-                            not_as_alias_handler(node.names),
-                            as_alias_handler(node.names)
-                        )
-                    return self.add_module(
+        # not relative
+        for module in self.local_modules:
+            if node.module == module[0]:
+                if os.path.isdir(module[1]):
+                    return self.from_directory_import(
                         module,
-                        None,
-                        as_alias_handler(node.names),
-                        retrieve_import_alias_mapping(node.names),
-                        from_from=True
+                        not_as_alias_handler(node.names),
+                        as_alias_handler(node.names)
                     )
-            for module in self.project_modules:
-                name = module[0]
-                if node.module == name:
-                    if os.path.isdir(module[1]):
-                        return self.from_directory_import(
-                            module,
-                            not_as_alias_handler(node.names),
-                            as_alias_handler(node.names),
-                            retrieve_import_alias_mapping(node.names)
-                        )
-                    return self.add_module(
+                return self.add_module(
+                    module=module,
+                    module_or_package_name=None,
+                    local_names=as_alias_handler(node.names),
+                    import_alias_mapping=retrieve_import_alias_mapping(node.names),
+                    from_from=True
+                )
+        for module in self.project_modules:
+            name = module[0]
+            if node.module == name:
+                if os.path.isdir(module[1]):
+                    return self.from_directory_import(
                         module,
-                        None,
+                        not_as_alias_handler(node.names),
                         as_alias_handler(node.names),
-                        retrieve_import_alias_mapping(node.names),
-                        from_from=True
+                        retrieve_import_alias_mapping(node.names)
                     )
+                return self.add_module(
+                    module=module,
+                    module_or_package_name=None,
+                    local_names=as_alias_handler(node.names),
+                    import_alias_mapping=retrieve_import_alias_mapping(node.names),
+                    from_from=True
+                )
 
-        if node.module in uninspectable_modules:
-            # Remember aliases for blackboxed and built-in imports such that we can label them fully qualified
-            # e.g. we want a call to "os.system" be recognised, even if we do "from os import system"
-            # from os import system as mysystem -> module=os, name=system, asname=mysystem
-            for name in node.names:
-                self.bb_or_bi_aliases[name.asname or name.name] = "{}.{}".format(node.module, name.name)
-        else:
-            log.warn("Cannot inspect module %s", node.module)
+        # Remember aliases for uninspecatble modules such that we can label them fully qualified
+        # e.g. we want a call to "os.system" be recognised, even if we do "from os import system"
+        # from os import system as mysystem -> module=os, name=system, asname=mysystem
+        for name in node.names:
+            local_definitions = self.module_definitions_stack[-1]
+            local_definitions.import_alias_mapping[name.asname or name.name] = "{}.{}".format(node.module, name.name)
+
+        if node.module not in uninspectable_modules:
+            log.warning("Cannot inspect module %s", node.module)
             uninspectable_modules.add(node.module)
         return IgnoredNode()

tests/vulnerabilities/vulnerabilities_across_files_test.py

@@ -62,7 +62,7 @@ def test_blackbox_library_call(self):
         EXPECTED_VULNERABILITY_DESCRIPTION = """
             File: examples/vulnerable_code_across_files/blackbox_library_call.py
              > User input at line 12, source "request.args.get(":
-                ~call_1 = ret_request.args.get('suggestion')
+                ~call_1 = ret_flask.request.args.get('suggestion')
             Reassigned in:
                 File: examples/vulnerable_code_across_files/blackbox_library_call.py
                  > Line 12: param = ~call_1

tests/vulnerabilities/vulnerabilities_base_test_case.py

@@ -11,7 +11,7 @@ def string_compare_alpha(self, output, expected_string):
 
     def assertAlphaEqual(self, output, expected_string):
         self.assertEqual(
-            [char for char in output if char.isalpha()],
-            [char for char in expected_string if char.isalpha()]
+            ''.join(char for char in output if char.isalpha()),
+            ''.join(char for char in expected_string if char.isalpha())
         )
         return True