Incremental changes mostly moving command parser out of victim

pwncat/__init__.py

@@ -1,8 +1,89 @@
 #!/usr/bin/env python3
 from typing import Optional
+from io import TextIOWrapper
+import sys
+import os
 
-from .config import Config
+import selectors
+from sqlalchemy.exc import InvalidRequestError
 
+# These need to be assigned prior to importing other
+# parts of pwncat
 victim: Optional["pwncat.remote.Victim"] = None
 
+from .config import Config
+from .commands import parser
+from .util import console
+from .db import get_session
+
 config: Config = Config()
+
+
+def interactive(platform):
+    """ Run the interactive pwncat shell with the given initialized victim.
+    This function handles the pwncat and remote prompts and does not return
+    until explicitly exited by the user.
+
+    This doesn't work yet. It's dependant on the new platform and channel
+    interface that isn't working yet, but it's what I'd like the eventual
+    interface to look like.
+
+    :param platform: an initialized platform object with a valid channel
+    :type platform: pwncat.platform.Platform
+    """
+
+    global victim
+    global config
+
+    # Initialize a new victim
+    victim = platform
+
+    # Ensure the prompt is initialized
+    parser.setup_prompt()
+
+    # Ensure our stdin reference is unbuffered
+    sys.stdin = TextIOWrapper(
+        os.fdopen(sys.stdin.fileno(), "br", buffering=0),
+        write_through=True,
+        line_buffering=False,
+    )
+
+    # Ensure we are in raw mode
+    parser.raw_mode()
+
+    # Create selector for asynchronous IO
+    selector = selectors.DefaultSelector()
+    selector.register(sys.stdin, selectors.EVENT_READ, None)
+    selector.register(victim.channel, selectors.EVENT_READ, None)
+
+    # Main loop state
+    done = False
+
+    try:
+        while not done:
+
+            for key, _ in selector.select():
+                if key.fileobj is sys.stdin:
+                    data = sys.stdin.buffer.read(64)
+                    data = parser.parse_prefix(data)
+                    if data:
+                        victim.channel.send(data)
+                else:
+                    data = victim.channel.recv(4096)
+                    if data is None or not data:
+                        done = True
+                        break
+                    sys.stdout.buffer.write(data)
+                    sys.stdout.flush()
+    except ConnectionResetError:
+        console.log("[yellow]warning[/yellow]: connection reset by remote host")
+    except SystemExit:
+        console.log("closing connection")
+    finally:
+        # Ensure the terminal is back to normal
+        parser.restore_term()
+        try:
+            # Commit any pending changes to the database
+            get_session().commit()
+        except InvalidRequestError:
+            pass

pwncat/__main__.py

@@ -16,6 +16,7 @@ from paramiko.buffered_pipe import BufferedPipe
 import pwncat
 from pwncat.util import console
 from pwncat.remote import Victim
+from pwncat.db import get_session
 
 
 def main():
@@ -113,7 +114,7 @@ def main():
             pwncat.victim.restore_local_term()
             try:
                 # Make sure everything was committed
-                pwncat.victim.session.commit()
+                get_session().commit()
             except InvalidRequestError:
                 pass
 

pwncat/channel/reconnect.py

@@ -22,7 +22,7 @@ class Reconnect(Channel):
             host = "0.0.0.0"
 
         if port is None:
-            raise ChannelError(f"no port specified")
+            raise ChannelError("no port specified")
 
         with Progress(
             f"bound to [blue]{host}[/blue]:[cyan]{port}[/cyan]",

pwncat/commands/__init__.py

@@ -24,9 +24,14 @@ from prompt_toolkit.history import InMemoryHistory, History
 from typing import Dict, Any, List, Iterable
 from colorama import Fore
 from enum import Enum, auto
+from io import TextIOWrapper
 import argparse
 import pkgutil
 import shlex
+import sys
+import fcntl
+import termios
+import tty
 import os
 import re
 
@@ -36,6 +41,7 @@ import pwncat
 import pwncat.db
 from pwncat.commands.base import CommandDefinition, Complete
 from pwncat.util import State, console
+from pwncat.db import get_session
 
 
 def resolve_blocks(source: str):
@@ -101,7 +107,8 @@ class DatabaseHistory(History):
     def load_history_strings(self) -> Iterable[str]:
         """ Load the history from the database """
         for history in (
-            pwncat.victim.session.query(pwncat.db.History)
+            get_session()
+            .query(pwncat.db.History)
             .order_by(pwncat.db.History.id.desc())
             .all()
         ):
@@ -110,12 +117,15 @@ class DatabaseHistory(History):
     def store_string(self, string: str) -> None:
         """ Store a command in the database """
         history = pwncat.db.History(host_id=pwncat.victim.host.id, command=string)
-        pwncat.victim.session.add(history)
+        get_session().add(history)
 
 
 class CommandParser:
     """ Handles dynamically loading command classes, parsing input, and
-    dispatching commands. """
+    dispatching commands. This class effectively has complete control over
+    the terminal whenever in an interactive pwncat session. It will change
+    termios modes for the control tty at will in order to support raw vs
+    command mode. """
 
     def __init__(self):
         """ We need to dynamically load commands from pwncat.commands """
@@ -134,6 +144,10 @@ class CommandParser:
         self.loading_complete = False
         self.aliases: Dict[str, CommandDefinition] = {}
         self.shortcuts: Dict[str, CommandDefinition] = {}
+        self.found_prefix: bool = False
+        # Saved terminal state to support switching between raw and normal
+        # mode.
+        self.saved_term_state = None
 
     def setup_prompt(self):
         """ This needs to happen after __init__ when the database is fully
@@ -210,10 +224,7 @@ class CommandParser:
 
                 if pwncat.config.module:
                     self.prompt.message = [
-                        (
-                            "fg:ansiyellow bold",
-                            f"({pwncat.config.module.name}) ",
-                        ),
+                        ("fg:ansiyellow bold", f"({pwncat.config.module.name}) ",),
                         ("fg:ansimagenta bold", "pwncat"),
                         ("", "$ "),
                     ]
@@ -315,6 +326,113 @@ class CommandParser:
             # The arguments were incorrect
             return
 
+    def parse_prefix(self, channel, data: bytes):
+        """ Parse data received from the user when in pwncat's raw mode.
+        This will intercept key presses from the user and interpret the
+        prefix and any bound keyboard shortcuts. It also sends any data
+        without a prefix to the remote channel.
+
+        :param data: input data from user
+        :type data: bytes
+        """
+
+        buffer = b""
+
+        for c in data:
+            if not self.found_prefix and c != pwncat.config["prefix"].value:
+                buffer += c
+                continue
+            elif not self.found_prefix and c == pwncat.config["prefix"].value:
+                self.found_prefix = True
+                channel.send(buffer)
+                buffer = b""
+                continue
+            elif self.found_prefix:
+                try:
+                    binding = pwncat.config.binding(c)
+                    if binding.strip() == "pass":
+                        buffer += c
+                    else:
+                        # Restore the normal terminal
+                        self.restore_term()
+
+                        # Run the binding script
+                        self.eval(binding, "<binding>")
+
+                        # Drain any channel output
+                        channel.drain()
+                        channel.send(b"\n")
+
+                        # Go back to a raw terminal
+                        self.raw_mode()
+                except KeyError:
+                    pass
+                self.found_prefix = False
+
+        # Flush any remaining raw data bound for the victim
+        channel.send(buffer)
+
+    def raw_mode(self):
+        """ Save the current terminal state and enter raw mode.
+        If the terminal is already in raw mode, this function
+        does nothing. """
+
+        if self.saved_term_state is not None:
+            return
+
+        # Ensure we don't have any weird buffering issues
+        sys.stdout.flush()
+
+        # Python doesn't provide a way to use setvbuf, so we reopen stdout
+        # and specify no buffering. Duplicating stdin allows the user to press C-d
+        # at the local prompt, and still be able to return to the remote prompt.
+        try:
+            os.dup2(sys.stdin.fileno(), sys.stdout.fileno())
+        except OSError:
+            pass
+        sys.stdout = TextIOWrapper(
+            os.fdopen(os.dup(sys.stdin.fileno()), "bw", buffering=0),
+            write_through=True,
+            line_buffering=False,
+        )
+
+        # Grab and duplicate current attributes
+        fild = sys.stdin.fileno()
+        old = termios.tcgetattr(fild)
+        new = termios.tcgetattr(fild)
+
+        # Remove ECHO from lflag and ensure we won't block
+        new[3] &= ~(termios.ECHO | termios.ICANON)
+        new[6][termios.VMIN] = 0
+        new[6][termios.VTIME] = 0
+        termios.tcsetattr(fild, termios.TCSADRAIN, new)
+
+        # Set raw mode
+        tty.setraw(sys.stdin)
+
+        orig_fl = fcntl.fcntl(sys.stdin, fcntl.F_GETFL)
+        fcntl.fcntl(sys.stdin, fcntl.F_SETFL, orig_fl)
+
+        self.saved_term_state = old, orig_fl
+
+    def restore_term(self, new_line=True):
+        """ Restores the normal terminal settings. This does nothing if the
+        terminal is not currently in raw mode. """
+
+        if self.saved_term_state is None:
+            return
+
+        termios.tcsetattr(
+            sys.stdin.fileno(), termios.TCSADRAIN, self.saved_term_state[0]
+        )
+        # tty.setcbreak(sys.stdin)
+        fcntl.fcntl(sys.stdin, fcntl.F_SETFL, self.saved_term_state[1])
+
+        if new_line:
+            sys.stdout.write("\n")
+
+        self.saved_term_state = None
+
 
 class CommandLexer(RegexLexer):
 
@@ -537,3 +655,9 @@ class CommandCompleter(Completer):
             yield from next_completer.get_completions(document, complete_event)
         elif this_completer is not None:
             yield from this_completer.get_completions(document, complete_event)
+
+
+# Here, we allocate the global parser object and initialize in-memory
+# settings
+parser: CommandParser = CommandParser()
+parser.setup_prompt()

pwncat/commands/busybox.py

@@ -11,6 +11,7 @@ from pwncat.commands.base import (
     StoreForAction,
 )
 from pwncat.util import console
+from pwncat.db import get_session
 
 
 class Command(CommandDefinition):
@@ -71,9 +72,13 @@ class Command(CommandDefinition):
                 return
 
             # Find all binaries which are provided by busybox
-            provides = pwncat.victim.session.query(pwncat.db.Binary).filter(
-                pwncat.db.Binary.path.contains(pwncat.victim.host.busybox),
-                pwncat.db.Binary.host_id == pwncat.victim.host.id,
+            provides = (
+                get_session()
+                .query(pwncat.db.Binary)
+                .filter(
+                    pwncat.db.Binary.path.contains(pwncat.victim.host.busybox),
+                    pwncat.db.Binary.host_id == pwncat.victim.host.id,
+                )
             )
 
             for binary in provides:
@@ -88,7 +93,8 @@ class Command(CommandDefinition):
 
             # Find all binaries which are provided from busybox
             nprovides = (
-                pwncat.victim.session.query(pwncat.db.Binary)
+                get_session()
+                .query(pwncat.db.Binary)
                 .filter(
                     pwncat.db.Binary.path.contains(pwncat.victim.host.busybox),
                     pwncat.db.Binary.host_id == pwncat.victim.host.id,

pwncat/commands/connect.py

@@ -21,6 +21,7 @@ from pwncat.commands.base import (
 
 # from pwncat.persist import PersistenceError
 from pwncat.modules.persist import PersistError
+from pwncat.db import get_session
 
 
 class Command(CommandDefinition):
@@ -119,7 +120,7 @@ class Command(CommandDefinition):
             # persistence methods
             hosts = {
                 host.hash: (host, [])
-                for host in pwncat.victim.session.query(pwncat.db.Host).all()
+                for host in get_session().query(pwncat.db.Host).all()
             }
 
             for module in modules:
@@ -201,9 +202,7 @@ class Command(CommandDefinition):
             try:
                 addr = ipaddress.ip_address(socket.gethostbyname(host))
                 row = (
-                    pwncat.victim.session.query(pwncat.db.Host)
-                    .filter_by(ip=str(addr))
-                    .first()
+                    get_session().query(pwncat.db.Host).filter_by(ip=str(addr)).first()
                 )
                 if row is None:
                     console.log(f"{level}: {str(addr)}: not found in database")

pwncat/commands/set.py

@@ -6,15 +6,14 @@ from sqlalchemy.orm import sessionmaker
 import pwncat
 from pwncat.commands.base import CommandDefinition, Complete, Parameter
 from pwncat.util import console, State
+from pwncat.db import get_session, reset_engine
 
 
 class Command(CommandDefinition):
     """ Set variable runtime variable parameters for pwncat """
 
     def get_config_variables(self):
-        options = (
-            ["state"] + list(pwncat.config.values) + list(pwncat.victim.users)
-        )
+        options = ["state"] + list(pwncat.config.values) + list(pwncat.victim.users)
 
         if pwncat.config.module:
             options.extend(pwncat.config.module.ARGUMENTS.keys())
@@ -82,16 +81,14 @@ class Command(CommandDefinition):
                     if args.variable == "db":
                         # We handle this specially to ensure the database is available
                         # as soon as this config is set
-                        pwncat.victim.engine = create_engine(
-                            pwncat.config["db"], echo=False
-                        )
-                        pwncat.db.Base.metadata.create_all(pwncat.victim.engine)
-
-                        # Create the session_maker and default session
-                        pwncat.victim.session_maker = sessionmaker(
-                            bind=pwncat.victim.engine
-                        )
-                        pwncat.victim.session = pwncat.victim.session_maker()
+                        reset_engine()
+                        if pwncat.victim.host is not None:
+                            pwncat.victim.host = (
+                                get_session()
+                                .query(pwncat.db.Host)
+                                .filter_by(id=pwncat.victim.host.id)
+                                .scalar()
+                            )
                 except ValueError as exc:
                     console.log(f"[red]error[/red]: {exc}")
             elif args.variable is not None:

pwncat/config.py

@@ -9,7 +9,6 @@ from prompt_toolkit.input.ansi_escape_sequences import (
     ANSI_SEQUENCES,
 )
 from prompt_toolkit.keys import ALL_KEYS, Keys
-import commentjson as json
 
 from pwncat.modules import BaseModule
 

pwncat/db/__init__.py

@@ -1,5 +1,9 @@
 #!/usr/bin/env python3
 
+from sqlalchemy.engine import Engine, create_engine
+from sqlalchemy.orm import Session, sessionmaker
+
+import pwncat
 from pwncat.db.base import Base
 from pwncat.db.binary import Binary
 from pwncat.db.history import History
@@ -9,3 +13,53 @@ from pwncat.db.suid import SUID
 from pwncat.db.tamper import Tamper
 from pwncat.db.user import User, Group, SecondaryGroupAssociation
 from pwncat.db.fact import Fact
+
+ENGINE: Engine = None
+SESSION_MAKER = None
+SESSION: Session = None
+
+
+def get_engine() -> Engine:
+    """
+    Get a copy of the database engine
+    """
+
+    global ENGINE
+
+    if ENGINE is not None:
+        return ENGINE
+
+    ENGINE = create_engine(pwncat.config["db"], echo=False)
+    Base.metadata.create_all(ENGINE)
+
+    return ENGINE
+
+
+def get_session() -> Session:
+    """
+    Get a new session object
+    """
+
+    global SESSION_MAKER
+    global SESSION
+
+    if SESSION_MAKER is None:
+        SESSION_MAKER = sessionmaker(bind=get_engine())
+    if SESSION is None:
+        SESSION = SESSION_MAKER()
+
+    return SESSION
+
+
+def reset_engine():
+    """
+    Reload the engine and session
+    """
+
+    global ENGINE
+    global SESSION
+    global SESSION_MAKER
+
+    ENGINE = None
+    SESSION = None
+    SESSION_MAKER = None

pwncat/modules/enumerate/__init__.py

@@ -8,6 +8,7 @@ import sqlalchemy
 import pwncat
 from pwncat.platform import Platform
 from pwncat.modules import BaseModule, Status, Argument, List
+from pwncat.db import get_session
 
 
 class Schedule(Enum):
@@ -61,14 +62,17 @@ class EnumerateModule(BaseModule):
 
         if clear:
             # Delete enumerated facts
-            query = pwncat.victim.session.query(pwncat.db.Fact).filter_by(
-                source=self.name, host_id=pwncat.victim.host.id
+            query = (
+                get_session()
+                .query(pwncat.db.Fact)
+                .filter_by(source=self.name, host_id=pwncat.victim.host.id)
             )
             query.delete(synchronize_session=False)
             # Delete our marker
             if self.SCHEDULE != Schedule.ALWAYS:
                 query = (
-                    pwncat.victim.session.query(pwncat.db.Fact)
+                    get_session()
+                    .query(pwncat.db.Fact)
                     .filter_by(host_id=pwncat.victim.host.id, type="marker")
                     .filter(pwncat.db.Fact.source.startswith(self.name))
                 )
@@ -77,7 +81,8 @@ class EnumerateModule(BaseModule):
 
         # Yield all the know facts which have already been enumerated
         existing_facts = (
-            pwncat.victim.session.query(pwncat.db.Fact)
+            get_session()
+            .query(pwncat.db.Fact)
             .filter_by(source=self.name, host_id=pwncat.victim.host.id)
             .filter(pwncat.db.Fact.type != "marker")
         )
@@ -92,7 +97,8 @@ class EnumerateModule(BaseModule):
 
         if self.SCHEDULE != Schedule.ALWAYS:
             exists = (
-                pwncat.victim.session.query(pwncat.db.Fact.id)
+                get_session()
+                .query(pwncat.db.Fact.id)
                 .filter_by(
                     host_id=pwncat.victim.host.id, type="marker", source=marker_name
                 )
@@ -114,11 +120,11 @@ class EnumerateModule(BaseModule):
                 host_id=pwncat.victim.host.id, type=typ, data=data, source=self.name
             )
             try:
-                pwncat.victim.session.add(row)
+                get_session().add(row)
                 pwncat.victim.host.facts.append(row)
-                pwncat.victim.session.commit()
+                get_session().commit()
             except sqlalchemy.exc.IntegrityError:
-                pwncat.victim.session.rollback()
+                get_session().rollback()
                 yield Status(data)
                 continue
 
@@ -140,7 +146,7 @@ class EnumerateModule(BaseModule):
                 source=marker_name,
                 data=None,
             )
-            pwncat.victim.session.add(row)
+            get_session().add(row)
             pwncat.victim.host.facts.append(row)
 
     def enumerate(self):

pwncat/modules/enumerate/gather.py

@@ -13,6 +13,7 @@ import pwncat.modules
 from pwncat import util
 from pwncat.util import console
 from pwncat.modules.enumerate import EnumerateModule
+from pwncat.db import get_session
 
 
 def strip_markup(styled_text: str) -> str:
@@ -86,7 +87,7 @@ class Module(pwncat.modules.BaseModule):
             for module in modules:
                 yield pwncat.modules.Status(module.name)
                 module.run(progress=self.progress, clear=True)
-            pwncat.victim.session.commit()
+            get_session().commit()
             pwncat.victim.reload_host()
             return
 

pwncat/modules/persist/__init__.py

@@ -15,6 +15,7 @@ from pwncat.modules import (
     PersistError,
     PersistType,
 )
+from pwncat.db import get_session
 
 
 class PersistModule(BaseModule):
@@ -95,7 +96,8 @@ class PersistModule(BaseModule):
 
         # Check if this module has been installed with the same arguments before
         ident = (
-            pwncat.victim.session.query(pwncat.db.Persistence.id)
+            get_session()
+            .query(pwncat.db.Persistence.id)
             .filter_by(host_id=pwncat.victim.host.id, method=self.name, args=kwargs)
             .scalar()
         )
@@ -110,7 +112,7 @@ class PersistModule(BaseModule):
                 yield result
 
             # Remove from the database
-            pwncat.victim.session.query(pwncat.db.Persistence).filter_by(
+            get_session().query(pwncat.db.Persistence).filter_by(
                 host_id=pwncat.victim.host.id, method=self.name, args=kwargs
             ).delete(synchronize_session=False)
             return
@@ -178,7 +180,7 @@ class PersistModule(BaseModule):
         )
         pwncat.victim.host.persistence.append(row)
 
-        pwncat.victim.session.commit()
+        get_session().commit()
 
     def install(self, **kwargs):
         """

pwncat/modules/persist/gather.py

@@ -5,6 +5,7 @@ import pwncat
 from pwncat.util import console
 from pwncat.modules import BaseModule, Argument, Status, Bool, Result
 import pwncat.modules.persist
+from pwncat.db import get_session
 
 
 @dataclasses.dataclass
@@ -93,11 +94,13 @@ class Module(BaseModule):
         """ Execute this module """
 
         if pwncat.victim.host is not None:
-            query = pwncat.victim.session.query(pwncat.db.Persistence).filter_by(
-                host_id=pwncat.victim.host.id
+            query = (
+                get_session()
+                .query(pwncat.db.Persistence)
+                .filter_by(host_id=pwncat.victim.host.id)
             )
         else:
-            query = pwncat.victim.session.query(pwncat.db.Persistence)
+            query = get_session().query(pwncat.db.Persistence)
 
         if module is not None:
             query = query.filter_by(method=module)

pwncat/remote/victim.py

@@ -35,6 +35,7 @@ from pwncat.remote import RemoteService
 from pwncat.tamper import TamperManager
 from pwncat.util import State, console
 from pwncat.modules.persist import PersistError, PersistType
+from pwncat.db import get_session
 
 
 def remove_busybox_tamper():
@@ -135,29 +136,12 @@ class Victim:
         self.client: Optional[socket.SocketType] = None
         # The shell we are running under on the remote host
         self.shell: str = "unknown"
-        # Database engine
-        self.engine: Engine = None
-        # Database session
-        self.session: Session = None
         # The host object as seen by the database
         self.host: pwncat.db.Host = None
         # The current user. This is cached while at the `pwncat` prompt
         # and reloaded whenever returning from RAW mode.
         self.cached_user: str = None
 
-        # The db engine is created here, but likely wrong. This happens
-        # before a configuration script is loaded, so likely creates a
-        # in memory db. This needs to happen because other parts of the
-        # framework assume a db engine exists, and therefore needs this
-        # reference. Also, in the case a config isn't loaded this
-        # needs to happen.
-        self.engine = create_engine(pwncat.config["db"], echo=False)
-        pwncat.db.Base.metadata.create_all(self.engine)
-
-        # Create the session_maker and default session
-        self.session_maker = sessionmaker(bind=self.engine)
-        self.session = self.session_maker()
-
     def reconnect(
         self, hostid: str, requested_method: str = None, requested_user: str = None
     ):
@@ -176,17 +160,8 @@ class Victim:
             will be tried.
         """
 
-        # Create the database engine, and then create the schema
-        # if needed.
-        self.engine = create_engine(pwncat.config["db"], echo=False)
-        pwncat.db.Base.metadata.create_all(self.engine)
-
-        # Create the session_maker and default session
-        self.session_maker = sessionmaker(bind=self.engine)
-        self.session = self.session_maker()
-
         # Load this host from the database
-        self.host = self.session.query(pwncat.db.Host).filter_by(hash=hostid).first()
+        self.host = get_session().query(pwncat.db.Host).filter_by(hash=hostid).first()
         if self.host is None:
             raise PersistError(f"invalid host hash")
 
@@ -235,17 +210,6 @@ class Victim:
         :return: None
         """
 
-        # Create the database engine, and then create the schema
-        # if needed.
-        if self.engine is None:
-            self.engine = create_engine(pwncat.config["db"], echo=False)
-            pwncat.db.Base.metadata.create_all(self.engine)
-
-            # Create the session_maker and default session
-            if self.session is None:
-                self.session_maker = sessionmaker(bind=self.engine)
-                self.session = self.session_maker()
-
         # Initialize the socket connection
         self.client = client
 
@@ -315,7 +279,7 @@ class Victim:
 
             # Lookup the remote host in our database. If it's not there, create an entry
             self.host = (
-                self.session.query(pwncat.db.Host).filter_by(hash=host_hash).first()
+                get_session().query(pwncat.db.Host).filter_by(hash=host_hash).first()
             )
             if self.host is None:
                 progress.log(f"new host w/ hash [cyan]{host_hash}[/cyan]")
@@ -324,9 +288,9 @@ class Victim:
                 # Probe for system information
                 self.probe_host_details(progress, task_id)
                 # Add the host to the session
-                self.session.add(self.host)
+                get_session().add(self.host)
                 # Commit what we know
-                self.session.commit()
+                get_session().commit()
 
             # Save the remote host IP address
             self.host.ip = self.client.getpeername()[0]
@@ -554,16 +518,17 @@ class Victim:
         # Replace anything we provide in our binary cache with the busybox version
         for name in provides:
             binary = (
-                self.session.query(pwncat.db.Binary)
+                get_session()
+                .query(pwncat.db.Binary)
                 .filter_by(host_id=self.host.id, name=name)
                 .first()
             )
             if binary is not None:
-                self.session.delete(binary)
+                get_session().delete(binary)
             binary = pwncat.db.Binary(name=name, path=f"{busybox_remote_path} {name}")
             self.host.binaries.append(binary)
 
-        self.session.commit()
+        get_session().commit()
 
         console.log(f"busybox installed w/ {len(provides)} applets")
 
@@ -572,7 +537,7 @@ class Victim:
         operations such as clearing enumeration data. """
 
         self.host = (
-            self.session.query(pwncat.db.Host).filter_by(id=self.host.id).first()
+            get_session().query(pwncat.db.Host).filter_by(id=self.host.id).first()
         )
 
     def probe_host_details(self, progress: Progress, task_id):
@@ -633,7 +598,7 @@ class Victim:
 
         for binary in self.host.binaries:
             if self.host.busybox in binary.path:
-                self.session.delete(binary)
+                get_session().delete(binary)
 
         # Did we upload a copy of busybox or was it already installed?
         if self.host.busybox_uploaded:
@@ -663,7 +628,8 @@ class Victim:
         """
 
         binary = (
-            self.session.query(pwncat.db.Binary)
+            get_session()
+            .query(pwncat.db.Binary)
             .filter_by(name=name, host_id=self.host.id)
             .first()
         )
@@ -2067,7 +2033,8 @@ class Victim:
                     continue
                 line = line.strip().split(":")
                 user = (
-                    self.session.query(pwncat.db.User)
+                    get_session()
+                    .query(pwncat.db.User)
                     .filter_by(host_id=self.host.id, id=int(line[2]), name=line[0])
                     .first()
                 )
@@ -2086,7 +2053,7 @@ class Victim:
         # Remove users that don't exist anymore
         for user in self.host.users:
             if user.name not in current_users:
-                self.session.delete(user)
+                get_session().delete(user)
                 self.host.users.remove(user)
 
         with self.open("/etc/group", "r") as filp:
@@ -2097,7 +2064,8 @@ class Victim:
 
                 line = line.split(":")
                 group = (
-                    self.session.query(pwncat.db.Group)
+                    get_session()
+                    .query(pwncat.db.Group)
                     .filter_by(host_id=self.host.id, id=int(line[2]))
                     .first()
                 )
@@ -2111,7 +2079,8 @@ class Victim:
 
                 for username in line[3].split(","):
                     user = (
-                        self.session.query(pwncat.db.User)
+                        get_session()
+                        .query(pwncat.db.User)
                         .filter_by(host_id=self.host.id, name=username)
                         .first()
                     )
@@ -2131,7 +2100,8 @@ class Victim:
                         continue
 
                     user = (
-                        self.session.query(pwncat.db.User)
+                        get_session()
+                        .query(pwncat.db.User)
                         .filter_by(host_id=self.host.id, name=entries[0])
                         .first()
                     )
@@ -2146,7 +2116,7 @@ class Victim:
 
         # Reload the host object
         self.host = (
-            self.session.query(pwncat.db.Host).filter_by(id=self.host.id).first()
+            get_session().query(pwncat.db.Host).filter_by(id=self.host.id).first()
         )
 
         return self.users

pwncat/tamper.py

@@ -6,6 +6,7 @@ from colorama import Fore
 
 import pwncat
 from pwncat.util import Access
+from pwncat.db import get_session
 
 
 class Action(Enum):
@@ -113,9 +114,9 @@ class LambdaTamper(Tamper):
 
 class TamperManager:
     """ TamperManager not only provides some automated ability to tamper with
-    properties of the remote system, but also a tracker for all modifications 
+    properties of the remote system, but also a tracker for all modifications
     on the remote system with the ability to remove previous changes. Other modules
-    can register system changes with `PtyHandler.tamper` in order to allow the 
+    can register system changes with `PtyHandler.tamper` in order to allow the
     user to get a wholistic view of all modifications of the remote system, and
     attempt revert all modifications automatically. """
 
@@ -147,7 +148,7 @@ class TamperManager:
         serialized = pickle.dumps(tamper)
         tracker = pwncat.db.Tamper(name=str(tamper), data=serialized)
         pwncat.victim.host.tampers.append(tracker)
-        pwncat.victim.session.commit()
+        get_session().commit()
 
     def custom(self, name: str, revert: Optional[Callable] = None):
         tamper = LambdaTamper(name, revert)
@@ -176,10 +177,8 @@ class TamperManager:
         It removes the tracking for this tamper. """
 
         tracker = (
-            pwncat.victim.session.query(pwncat.db.Tamper)
-            .filter_by(name=str(tamper))
-            .first()
+            get_session().query(pwncat.db.Tamper).filter_by(name=str(tamper)).first()
         )
         if tracker is not None:
-            pwncat.victim.session.delete(tracker)
-            pwncat.victim.session.commit()
+            get_session().delete(tracker)
+            get_session().commit()