Ruff enhancements (#697)

* change ruff config * replace isort with ruff * pass pyflakes * pass pycodestyle errors * run W293 * run W291 * run pycodestyle errors * run UP010 * run UP029 * run UP004 * run UP015 * run UP024 * run UP030 * run UP032 * run UP012 * run E721 * run UP034 * run UP031 * changelog and contributing.md * run black * set pipefail so | tee doesn't suppress the error * auto-fix bugbear * B028 * B006 * B904 * B005 * update contributing * fix launcher tests
BC-SECURITY · Sep 27, 2023 · 02a0496 · 02a0496
1 parent 4c44216
commit 02a0496
Show file tree

Hide file tree

Showing 226 changed files with 878 additions and 1,525 deletions.
diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md
@@ -46,18 +46,14 @@ The `main` branch in `BC-SECURITY/Empire` automatically syncs.
 ### Code Formatting and Linting
 
 * We are using [psf/black](https://github.com/psf/black) for code formatting.
-    * Black is a Python code formatter that helps to keep the codebase uniform and easy to read
-* We are using [PyCQA/isort](https://github.com/PyCQA/isort)
-    * Isort is a Python utility that sorts and formats imports.
 * We are using [charliermarsh/ruff](https://github.com/charliermarsh/ruff) for linting.
-    * Ruff is a python linter that helps identify common bugs and style issues.
+    * We are using the E, W, F, I, UP, and B rulesets.
 * After implementing your changes:
-    1. run `black .` (or `poetry run black .`).
-    2. run `isort .` (or `poetry run isort .`).
-    3. run `ruff . --fix` (or `poetry run ruff . --fix`). 
+    1. run `ruff . --fix` (or `poetry run ruff . --fix`).
+    2. run `black .` (or `poetry run black .`).
 * The repo is also configured to use [pre-commit](https://pre-commit.com/) to automatically format code.
   * Once you have pre-commit installed, you can run `pre-commit install` to install the pre-commit hooks.
-  * Then pre-commit will execute black, isort, and ruff automatically before committing.
+  * Then pre-commit will execute black and ruff automatically before committing.
 
 ### Tests
 

diff --git a/.github/workflows/lint-and-test.yml b/.github/workflows/lint-and-test.yml
@@ -13,9 +13,6 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - uses: psf/black@23.9.1
-      - uses: isort/isort-action@master
-        with:
-          isort-version: 5.12.0
       - name: Run ruff
         run: |
           pip install ruff==0.0.283
@@ -76,6 +73,7 @@ jobs:
           poetry install
       - name: Run test suite - mysql
         run: |
+            set -o pipefail
             if [ "${{ matrix.python-version }}" = "3.11" ]; then
                 DATABASE_USE=mysql poetry run pytest -v --runslow --cov=empire/server --junitxml=pytest.xml --cov-report=term-missing:skip-covered . | tee pytest-coverage.txt
             else

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -7,20 +7,15 @@ repos:
     - id: check-yaml
     - id: check-merge-conflict
     - id: end-of-file-fixer
--   repo: https://github.com/psf/black-pre-commit-mirror
-    rev: 23.9.1
-    hooks:
-    - id: black
-      language_version: python3.9
-
--   repo: https://github.com/pycqa/isort
-    rev: 5.12.0
-    hooks:
-      - id: isort
-        name: isort (python)
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
   rev: v0.0.283
   hooks:
     - id: ruff
       args: [--fix, --exit-non-zero-on-fix]
+
+-   repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 23.9.1
+    hooks:
+    - id: black
+      language_version: python3.9
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,8 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
--  Updated Dropbox C2 to use new API endpoints (@Cx01N)
--  Standardized Kill Date and Working Hours for PowerShell Agents (@Cx01N)
+-   Updated Dropbox C2 to use new API endpoints (@Cx01N)
+-   Standardized Kill Date and Working Hours for PowerShell Agents (@Cx01N)
+-   Apply fixes for future Python 3.12 compatibility (@Vinnybod)
+-   Add additional rulesets to ruff linting (@Vinnybod)
 
 ## [5.7.1] - 2023-09-25
 

diff --git a/empire/client/client.py b/empire/client/client.py
@@ -118,7 +118,7 @@ class CliExitException(BaseException):
     pass
 
 
-class EmpireCli(object):
+class EmpireCli:
     def __init__(self) -> None:
         self.completer = MyCustomCompleter(self)
         self.menus: Dict[Menu] = {
@@ -443,7 +443,7 @@ def parse_command_line(self, text: str, cmd_line: List[str], resource_file=False
                     args = self.strip(docopt(func.__doc__, argv=cmd_line[1:]))
                     new_args = {}
                     # todo casting for type hinted values?
-                    for key, hint in get_type_hints(func).items():
+                    for key in get_type_hints(func).keys():
                         if key != "return":
                             new_args[key] = args[key]
                     func(**new_args)

diff --git a/empire/client/src/EmpireCliConfig.py b/empire/client/src/EmpireCliConfig.py
@@ -7,7 +7,7 @@
 log = logging.getLogger(__name__)
 
 
-class EmpireCliConfig(object):
+class EmpireCliConfig:
     def __init__(self):
         self.yaml: Dict = {}
         if "--config" in sys.argv:
@@ -20,7 +20,7 @@ def __init__(self):
 
     def set_yaml(self, location: str):
         try:
-            with open(location, "r") as stream:
+            with open(location) as stream:
                 self.yaml = yaml.safe_load(stream)
         except yaml.YAMLError as exc:
             log.error(exc)

diff --git a/empire/client/src/EmpireCliState.py b/empire/client/src/EmpireCliState.py
@@ -24,7 +24,7 @@
     pass
 
 
-class EmpireCliState(object):
+class EmpireCliState:
     def __init__(self):
         self.host = ""
         self.port = ""
@@ -427,7 +427,7 @@ def get_agents(self):
         self.agents = {x["name"]: x for x in response.json()["records"]}
 
         # Whenever agents are refreshed, add socketio listeners for taskings.
-        for name, agent in self.agents.items():
+        for _name, agent in self.agents.items():
             session_id = agent["session_id"]
             self.sio.on(f"agents/{session_id}/task", self.add_to_cached_results)
 
@@ -664,7 +664,7 @@ def get_active_plugins(self):
         )
 
         self.plugins = {x["name"]: x for x in response.json()["records"]}
-        for name, plugin in self.plugins.items():
+        for _name, plugin in self.plugins.items():
             plugin_name = plugin["name"]
             self.sio.on(f"plugins/{plugin_name}/notifications", self.add_plugin_cache)
         return self.plugins

diff --git a/empire/client/src/Shortcut.py b/empire/client/src/Shortcut.py
@@ -7,7 +7,7 @@
 
 
 # https://yzhong-cs.medium.com/serialize-and-deserialize-complex-json-in-python-205ecc636caa
-class ShortcutParam(object):
+class ShortcutParam:
     def __init__(self, name: str, dynamic: bool = False, value: Optional[str] = ""):
         self.name = name
         self.dynamic = dynamic
@@ -18,7 +18,7 @@ def from_json(cls, data):
         return cls(**data)
 
 
-class Shortcut(object):
+class Shortcut:
     def __init__(
         self,
         name: str,

diff --git a/empire/client/src/bindings.py b/empire/client/src/bindings.py
@@ -43,6 +43,6 @@ def do_konami(event):
 
     color_depth = ColorDepth.DEPTH_8_BIT
     with ProgressBar(formatters=custom_formatters, color_depth=color_depth) as pb:
-        for i in pb(range(1000), label=""):
+        for _i in pb(range(1000), label=""):
             time.sleep(0.001)
     print("Downloaded L33t Hax...")
diff --git a/empire/client/src/menus/AdminMenu.py b/empire/client/src/menus/AdminMenu.py
@@ -187,7 +187,7 @@ def load_malleable_profile(
 
         Usage: load_malleable_profile <profile_directory> [profile_category]
         """
-        with open(profile_directory, "r") as stream:
+        with open(profile_directory) as stream:
             profile_data = stream.read()
 
         post_body = {

diff --git a/empire/client/src/menus/EditListenerMenu.py b/empire/client/src/menus/EditListenerMenu.py
@@ -74,7 +74,7 @@ def execute(self):
         # Hopefully this will force us to provide more info in api errors ;)
         post_body = {}
         temp_record = {}
-        for key, value in self.record_options.items():
+        for key in self.record_options.keys():
             post_body[key] = self.record_options[key]["value"]
 
         temp_record["options"] = post_body

diff --git a/empire/client/src/menus/InteractMenu.py b/empire/client/src/menus/InteractMenu.py
@@ -495,7 +495,7 @@ def execute_shortcut(self, command_name: str, params: List[str]):
 
         # TODO Still haven't figured out other data types. Right now everything is a string.
         #  Which I think is how it is in the old cli
-        for key, value in module_options.items():
+        for key in module_options.keys():
             if key in shortcut.get_dynamic_param_names():
                 # Grab filename, send to server, and save a copy off in the downloads folder
                 if key in ["File"]:
@@ -569,7 +569,7 @@ def vnc(self) -> None:
         post_body = {}
         post_body["options"] = {}
 
-        for key, value in module_options.items():
+        for key in module_options.keys():
             post_body["options"][key] = str(module_options[key]["value"])
 
         post_body["module_id"] = "csharp_vnc_vncserver"

diff --git a/empire/client/src/menus/Menu.py b/empire/client/src/menus/Menu.py
@@ -8,7 +8,7 @@
 from empire.client.src.utils.cli_util import command
 
 
-class Menu(object):
+class Menu:
     """
     Base Menu object.
     """

diff --git a/empire/client/src/menus/ShellMenu.py b/empire/client/src/menus/ShellMenu.py
@@ -59,13 +59,17 @@ def update_directory(self, session_id: str):
         Update current directory
         """
         if self.language == "powershell":
-            task_id: int = state.agent_shell(session_id, "(Resolve-Path .\).Path")["id"]
+            task_id: int = state.agent_shell(session_id, r"(Resolve-Path .\).Path")[
+                "id"
+            ]
         elif self.language == "python":
             task_id: int = state.agent_shell(session_id, "echo $PWD")["id"]
         elif self.language == "ironpython":
             task_id: int = state.agent_shell(session_id, "cd .")["id"]
         elif self.language == "csharp":
-            task_id: int = state.agent_shell(session_id, "(Resolve-Path .\).Path")["id"]
+            task_id: int = state.agent_shell(session_id, r"(Resolve-Path .\).Path")[
+                "id"
+            ]
             pass
 
         count = 0

diff --git a/empire/client/src/menus/UseListenerMenu.py b/empire/client/src/menus/UseListenerMenu.py
@@ -68,7 +68,7 @@ def execute(self):
         # Hopefully this will force us to provide more info in api errors ;)
         post_body = {}
         temp_record = {}
-        for key, value in self.record_options.items():
+        for key in self.record_options.keys():
             post_body[key] = self.record_options[key]["value"]
 
         temp_record["options"] = post_body

diff --git a/empire/client/src/menus/UseMenu.py b/empire/client/src/menus/UseMenu.py
@@ -134,7 +134,7 @@ def set(self, key: str, value: str):
             value = value[1:-1]
         if key in self.record_options:
             self.record_options[key]["value"] = value
-            log.info("Set %s to %s" % (key, value))
+            log.info(f"Set {key} to {value}")
         else:
             log.error(f"Could not find field: {key}")
 

diff --git a/empire/client/src/menus/UseModuleMenu.py b/empire/client/src/menus/UseModuleMenu.py
@@ -104,7 +104,7 @@ def execute(self):
 
         post_body = {"options": {}}
 
-        for key, value in self.record_options.items():
+        for key in self.record_options.keys():
             post_body["options"][key] = self.record_options[key]["value"]
 
         post_body["module_id"] = self.record["id"]

diff --git a/empire/client/src/menus/UsePluginMenu.py b/empire/client/src/menus/UsePluginMenu.py
@@ -51,7 +51,7 @@ def display_cached_results(self) -> None:
         Print the plugin results for all the results that have been received for this plugin.
         """
         plugin_results = state.cached_plugin_results.get(self.selected, {})
-        for key, value in plugin_results.items():
+        for value in plugin_results.values():
             print(print_util.color(value))
 
         state.cached_plugin_results.get(self.selected, {}).clear()
@@ -77,7 +77,7 @@ def execute(self):
         """
         post_body = {}
         post_body["options"] = {}
-        for key, value in self.record_options.items():
+        for key in self.record_options.keys():
             post_body["options"][key] = self.record_options[key]["value"]
 
         response = state.execute_plugin(self.record["id"], post_body)

diff --git a/empire/client/src/menus/UseStagerMenu.py b/empire/client/src/menus/UseStagerMenu.py
@@ -85,7 +85,7 @@ def execute(self):
         # Hopefully this will force us to provide more info in api errors ;)
         post_body = {}
         temp_record = {}
-        for key, value in self.record_options.items():
+        for key in self.record_options.keys():
             post_body[key] = self.record_options[key]["value"]
 
         temp_record["options"] = post_body

diff --git a/empire/client/src/utils/print_util.py b/empire/client/src/utils/print_util.py
@@ -29,24 +29,24 @@ def color(string_name, color_name=None):
             str_list = string_name.split("\n")
             str_list_modified = []
             for s in str_list:
-                str_list_modified.append("\x1b[%sm%s\x1b[0m" % (";".join(attr), s))
+                str_list_modified.append("\x1b[{}m{}\x1b[0m".format(";".join(attr), s))
             return "\n".join(str_list_modified)
         else:
-            return "\x1b[%sm%s\x1b[0m" % (";".join(attr), string_name)
+            return "\x1b[{}m{}\x1b[0m".format(";".join(attr), string_name)
 
     else:
         if string_name.strip().startswith("[!]"):
             attr.append("31")
-            return "\x1b[%sm%s\x1b[0m" % (";".join(attr), string_name)
+            return "\x1b[{}m{}\x1b[0m".format(";".join(attr), string_name)
         elif string_name.strip().startswith("[+]"):
             attr.append("32")
-            return "\x1b[%sm%s\x1b[0m" % (";".join(attr), string_name)
+            return "\x1b[{}m{}\x1b[0m".format(";".join(attr), string_name)
         elif string_name.strip().startswith("[*]"):
             attr.append("34")
-            return "\x1b[%sm%s\x1b[0m" % (";".join(attr), string_name)
+            return "\x1b[{}m{}\x1b[0m".format(";".join(attr), string_name)
         elif string_name.strip().startswith("[>]"):
             attr.append("33")
-            return "\x1b[%sm%s\x1b[0m" % (";".join(attr), string_name)
+            return "\x1b[{}m{}\x1b[0m".format(";".join(attr), string_name)
         else:
             return string_name
 
@@ -84,11 +84,11 @@ def title(version, server, modules, listeners, agents):
     )
     print(
         """
-                   ███████╗███╗   ███╗██████╗ ██╗██████╗    ███████╗                                         
-                   ██╔════╝████╗ ████║██╔══██╗██║██╔══██╗   ██╔════╝                                         
-                   █████╗  ██╔████╔██║██████╔╝██║██████╔╝   █████╗                                          
-                   ██╔══╝  ██║╚██╔╝██║██╔═══╝ ██║██╔══██╗   ██╔══╝                                          
-                   ███████╗██║ ╚═╝ ██║██║     ██║██║  █████║███████╗  
+                   ███████╗███╗   ███╗██████╗ ██╗██████╗    ███████╗
+                   ██╔════╝████╗ ████║██╔══██╗██║██╔══██╗   ██╔════╝
+                   █████╗  ██╔████╔██║██████╔╝██║██████╔╝   █████╗
+                   ██╔══╝  ██║╚██╔╝██║██╔═══╝ ██║██╔══██╗   ██╔══╝
+                   ███████╗██║ ╚═╝ ██║██║     ██║██║  █████║███████╗
                    ╚══════╝╚═╝     ╚═╝╚═╝     ╚═╝╚═╝  ╚════╝╚══════╝
 
 """
@@ -152,13 +152,13 @@ def loading():
                            `-/osyhdddddhyo:
                                 ``.----.`
                         \x1b[0m
-                        
-          ███████╗███╗   ███╗██████╗ ██╗██████╗    ███████╗                                          
-          ██╔════╝████╗ ████║██╔══██╗██║██╔══██╗   ██╔════╝                                          
-          █████╗  ██╔████╔██║██████╔╝██║██████╔╝   █████╗                                          
-          ██╔══╝  ██║╚██╔╝██║██╔═══╝ ██║██╔══██╗   ██╔══╝                                         
-          ███████╗██║ ╚═╝ ██║██║     ██║██║  █████║███████╗    
-          ╚══════╝╚═╝     ╚═╝╚═╝     ╚═╝╚═╝  ╚════╝╚══════╝             
+
+          ███████╗███╗   ███╗██████╗ ██╗██████╗    ███████╗
+          ██╔════╝████╗ ████║██╔══██╗██║██╔══██╗   ██╔════╝
+          █████╗  ██╔████╔██║██████╔╝██║██████╔╝   █████╗
+          ██╔══╝  ██║╚██╔╝██║██╔═══╝ ██║██╔══██╗   ██╔══╝
+          ███████╗██║ ╚═╝ ██║██║     ██║██║  █████║███████╗
+          ╚══════╝╚═╝     ╚═╝╚═╝     ╚═╝╚═╝  ╚════╝╚══════╝
                         """
     )
 

diff --git a/empire/server/api/app.py b/empire/server/api/app.py
@@ -19,7 +19,7 @@
 log = logging.getLogger(__name__)
 
 
-class MyJsonWrapper(object):
+class MyJsonWrapper:
     @staticmethod
     def dumps(*args, **kwargs):
         if "cls" not in kwargs:

diff --git a/empire/server/api/jwt_auth.py b/empire/server/api/jwt_auth.py
@@ -83,8 +83,8 @@ async def get_current_user(
         if username is None:
             raise credentials_exception
         token_data = TokenData(username=username)
-    except JWTError:
-        raise credentials_exception
+    except JWTError as e:
+        raise credentials_exception from e
     user = get_user(db, username=token_data.username)
     if user is None:
         raise credentials_exception

diff --git a/empire/server/api/v2/download/download_dto.py b/empire/server/api/v2/download/download_dto.py
@@ -5,13 +5,7 @@
 from pydantic import BaseModel
 
 from empire.server.api.v2.tag.tag_dto import Tag, domain_to_dto_tag
-
-
-def removeprefix(value: str, prefix: str) -> str:
-    if value.startswith(prefix):
-        return value[len(prefix) :]
-    else:
-        return value[:]
+from empire.server.utils.data_util import removeprefix
 
 
 def domain_to_dto_download(download):