diff --git a/pwncat/__init__.py b/pwncat/__init__.py index a6317d5..7b47343 100644 --- a/pwncat/__init__.py +++ b/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 diff --git a/pwncat/__main__.py b/pwncat/__main__.py index 4dbd4c7..54694ae 100644 --- a/pwncat/__main__.py +++ b/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 diff --git a/pwncat/channel/reconnect.py b/pwncat/channel/reconnect.py index 1ac9702..6ab3a20 100644 --- a/pwncat/channel/reconnect.py +++ b/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]", diff --git a/pwncat/commands/__init__.py b/pwncat/commands/__init__.py index 806624a..87a38ff 100644 --- a/pwncat/commands/__init__.py +++ b/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, "") + + # 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() diff --git a/pwncat/commands/busybox.py b/pwncat/commands/busybox.py index 4599dc7..fe8ed80 100644 --- a/pwncat/commands/busybox.py +++ b/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, diff --git a/pwncat/commands/connect.py b/pwncat/commands/connect.py index 64bc8d7..64ae9fb 100644 --- a/pwncat/commands/connect.py +++ b/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") diff --git a/pwncat/commands/set.py b/pwncat/commands/set.py index ff4f7d3..739dc68 100644 --- a/pwncat/commands/set.py +++ b/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: diff --git a/pwncat/config.py b/pwncat/config.py index 5737cf8..6bdbcc3 100644 --- a/pwncat/config.py +++ b/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 diff --git a/pwncat/db/__init__.py b/pwncat/db/__init__.py index a1979bc..b59ec35 100644 --- a/pwncat/db/__init__.py +++ b/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 diff --git a/pwncat/modules/enumerate/__init__.py b/pwncat/modules/enumerate/__init__.py index 253120a..dc7f915 100644 --- a/pwncat/modules/enumerate/__init__.py +++ b/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): diff --git a/pwncat/modules/enumerate/gather.py b/pwncat/modules/enumerate/gather.py index ee79cc6..ff4d1e0 100644 --- a/pwncat/modules/enumerate/gather.py +++ b/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 diff --git a/pwncat/modules/persist/__init__.py b/pwncat/modules/persist/__init__.py index bf565d1..942cc3b 100644 --- a/pwncat/modules/persist/__init__.py +++ b/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): """ diff --git a/pwncat/modules/persist/gather.py b/pwncat/modules/persist/gather.py index 8d9688c..4a724be 100644 --- a/pwncat/modules/persist/gather.py +++ b/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) diff --git a/pwncat/remote/victim.py b/pwncat/remote/victim.py index 6df67c7..421114d 100644 --- a/pwncat/remote/victim.py +++ b/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 diff --git a/pwncat/tamper.py b/pwncat/tamper.py index 8b96c53..e1cd7ba 100644 --- a/pwncat/tamper.py +++ b/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()