Merge branch 'master' into issue-134-dockerhub-builds-failing

CHANGELOG.md

@@ -10,6 +10,10 @@ and simply didn't have the time to go back and retroactively create one.
 ## [Unreleased]
 ### Fixed
 - Pinned container base image to alpine 3.13.5 and installed to virtualenv ([#134](https://github.com/calebstewart/pwncat/issues/134))
+- Fixed syntax for f-strings in escalation command
+### Changed
+- Changed session tracking so session IDs aren't reused
+- Changed zsh prompt to match CWD of other shell prompts
 
 ## [0.4.2] - 2021-06-15
 Quick patch release due to corrected bug in `ChannelFile` which caused command

pwncat/__main__.py

@@ -290,9 +290,16 @@ def main():
                 transient=True,
             ) as progress:
                 task = progress.add_task("task", status="...")
-                while manager.sessions:
-                    progress.update(task, status=str(manager.sessions[0].platform))
-                    manager.sessions[0].close()
+
+                # Retrieve the existing session IDs list
+                session_ids = list(manager.sessions.keys())
+
+                # Close each session based on its ``session_id``
+                for session_id in session_ids:
+                    progress.update(
+                        task, status=str(manager.sessions[session_id].platform)
+                    )
+                    manager.sessions[session_id].close()
 
                 progress.update(task, status="done!", completed=100)
 

pwncat/commands/escalate.py

@@ -32,7 +32,7 @@ class Link:
 
         if self.escalation.type == "escalate.spawn":
             self.result.log(
-                "leaving behind open session as [cyan]{self.old_session.current_user().name}[/cyan]"
+                f"leaving behind open session as [cyan]{self.old_session.current_user().name}[/cyan]"
             )
 
         self.old_session.manager.target = self.old_session
@@ -163,7 +163,7 @@ class Command(CommandDefinition):
                     continue
                 except IndexError:
                     manager.target.log(
-                        "[red]error[/red]: no working escalation paths found for {user.name}"
+                        f"[red]error[/red]: no working escalation paths found for {user.name}"
                     )
                     break
 

pwncat/commands/sessions.py

@@ -40,15 +40,15 @@ class Command(CommandDefinition):
         if args.list or (not args.kill and args.session_id is None):
             table = Table(title="Active Sessions", box=box.MINIMAL_DOUBLE_HEAD)
 
-            table.add_column("")
+            table.add_column("ID")
             table.add_column("User")
             table.add_column("Host ID")
             table.add_column("Platform")
             table.add_column("Type")
             table.add_column("Address")
 
-            for ident, session in enumerate(manager.sessions):
-                ident = str(ident)
+            for session_id, session in manager.sessions.items():
+                ident = str(session_id)
                 kwargs = {"style": ""}
                 if session is manager.target:
                     ident = "*" + ident
@@ -71,7 +71,8 @@ class Command(CommandDefinition):
             console.log("[red]error[/red]: no session id specified")
             return
 
-        if args.session_id < 0 or args.session_id >= len(manager.sessions):
+        # check if a session with the provided ``session_id`` exists or not
+        if args.session_id not in manager.sessions:
             console.log(f"[red]error[/red]: {args.session_id}: no such session!")
             return
 

pwncat/manager.py

@@ -58,6 +58,7 @@ class Session:
         channel: Optional[Channel] = None,
         **kwargs,
     ):
+        self.id = manager.session_id
         self.manager = manager
         self.background = None
         self._db_session = None
@@ -87,7 +88,7 @@ class Session:
             )
 
         # Register this session with the manager
-        self.manager.sessions.append(self)
+        self.manager.sessions[self.id] = self
         self.manager.target = self
 
         # Initialize the host reference
@@ -290,10 +291,10 @@ class Session:
 
     def died(self):
 
-        if self not in self.manager.sessions:
+        if self.id not in self.manager.sessions:
             return
 
-        self.manager.sessions.remove(self)
+        del self.manager.sessions[self.id]
 
         if self.manager.target == self:
             self.manager.target = None
@@ -343,7 +344,8 @@ class Manager:
 
     def __init__(self, config: str = None):
         self.config = Config()
-        self.sessions: List[Session] = []
+        self.session_id = 0  # start with 0-indexed session IDs
+        self.sessions: Dict[int, Session] = {}
         self.modules: Dict[str, pwncat.modules.BaseModule] = {}
         self._target = None
         self.parser = CommandParser(self)
@@ -415,8 +417,12 @@ class Manager:
     def __exit__(self, _, __, ___):
         """Ensure all sessions are closed"""
 
-        while self.sessions:
-            self.sessions[0].close()
+        # Retrieve the existing session IDs list
+        session_ids = list(self.sessions.keys())
+
+        # Close each session based on its ``session_id``
+        for session_id in session_ids:
+            self.sessions[session_id].close()
 
     def open_database(self):
         """Create the internal engine and session builder
@@ -503,7 +509,7 @@ class Manager:
 
     @target.setter
     def target(self, value: Session):
-        if value is not None and value not in self.sessions:
+        if value is not None and value not in self.sessions.values():
             raise ValueError("invalid target")
         self._target = value
 
@@ -623,6 +629,13 @@ class Manager:
         """
 
         session = Session(self, platform, channel, **kwargs)
+
+        # Increment the ``session_id`` variable upon adding a new session
+        # Session constructor will automatically grab the current
+        # ``session_id`` from the ``manager`` object passed as the first argument
+
+        self.session_id += 1
+
         return session
 
     def _process_input(self, data: bytes, has_prefix: bool):

pwncat/platform/linux.py

@@ -495,7 +495,7 @@ class Linux(Platform):
     PROMPTS = {
         "sh": """'$(command printf "(remote) $(whoami)@$(hostname):$PWD\\$ ")'""",
         "dash": """'$(command printf "(remote) $(whoami)@$(hostname):$PWD\\$ ")'""",
-        "zsh": """'%B%F{red}(remote) %B%F{yellow}%n@%M%B%F{reset}:%B%F{cyan}%(6~.%-1~/…/%4~.%5~)%B%(#.%b%F{white}#.%b%F{white}$)%b%F{reset} '""",
+        "zsh": """'%B%F{red}(remote) %B%F{yellow}%n@%M%B%F{reset}:%B%F{cyan}$PWD%B%(#.%b%F{white}#.%b%F{white}$)%b%F{reset} '""",
         "default": """'$(command printf "\\[\\033[01;31m\\](remote)\\[\\033[0m\\] \\[\\033[01;33m\\]$(whoami)@$(hostname)\\[\\033[0m\\]:\\[\\033[1;36m\\]$PWD\\[\\033[0m\\]\\$ ")'""",
     }