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