From 3bccda6de497e39c061af7858e6da27ab38167e1 Mon Sep 17 00:00:00 2001
From: Mitul16 <mitul.varshney.2@gmail.com>
Date: Wed, 16 Jun 2021 16:56:52 +0530
Subject: [PATCH] Modified manager->sessions to work as a Dictionary

NOTE: This is for linux! This change will add ID management for the sessions, so that each session gets a unique ID (for a single `pwncat` process) and killing a session will not affect other session IDs. This is useful when we are pivoting using `pwncat`, and removes the confusion with ID management since the current implementation uses a List and works with list indices instead of session IDs
---
 pwncat/__main__.py          | 11 ++++++++---
 pwncat/commands/sessions.py |  9 +++++----
 pwncat/manager.py           | 30 +++++++++++++++++++++++-------
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/pwncat/__main__.py b/pwncat/__main__.py
index dcd5fc9..e2e949a 100644
--- a/pwncat/__main__.py
+++ b/pwncat/__main__.py
@@ -281,9 +281,14 @@ 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)
 
diff --git a/pwncat/commands/sessions.py b/pwncat/commands/sessions.py
index 88f93c5..000d5e7 100644
--- a/pwncat/commands/sessions.py
+++ b/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
 
diff --git a/pwncat/manager.py b/pwncat/manager.py
index 9c57108..ab16c55 100644
--- a/pwncat/manager.py
+++ b/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,16 @@ class Manager:
         """
 
         session = Session(self, platform, channel, **kwargs)
+
+        # increment the ``session_id`` variable upon adding a new session
+        # this may cause issues when a session is not properly configured
+        # for example when the victim machine disconnects or terminates the parent process
+        # in that case, we have already incremented the ``session_id``
+        # thus the ``session_id`` for the bad session is no longer in use
+        # check other modules to prevent this from happening
+
+        self.session_id += 1
+
         return session
 
     def _process_input(self, data: bytes, has_prefix: bool):