From 3678e9fa6694fdef746c16cb5445b478fd9ee027 Mon Sep 17 00:00:00 2001 From: Caleb Stewart Date: Fri, 5 Jun 2020 21:32:24 -0400 Subject: [PATCH] Added the rich module rich provides better progress bars and log output and exception tracebacks. --- pwncat/commands/__init__.py | 9 ++-- pwncat/commands/download.py | 66 +++++++++++++++++++++-------- pwncat/commands/upload.py | 83 +++++++++++++++++++++++++++---------- pwncat/data/gtfobins.json | 3 +- pwncat/file.py | 25 ++++++----- pwncat/remote/victim.py | 6 +-- pwncat/util.py | 4 ++ requirements.txt | 1 + setup.py | 1 + 9 files changed, 141 insertions(+), 57 deletions(-) diff --git a/pwncat/commands/__init__.py b/pwncat/commands/__init__.py index 23d80ca..5c36e7f 100644 --- a/pwncat/commands/__init__.py +++ b/pwncat/commands/__init__.py @@ -35,7 +35,7 @@ from pprint import pprint import pwncat import pwncat.db from pwncat.commands.base import CommandDefinition, Complete -from pwncat.util import State +from pwncat.util import State, console from pwncat import util @@ -224,8 +224,11 @@ class CommandParser: # We have a connection! Go back to raw mode pwncat.victim.state = State.RAW self.running = False - except (Exception, KeyboardInterrupt) as exc: - traceback.print_exc() + except Exception: + console.print_exception(width=None) + continue + except KeyboardInterrupt: + console.log("Keyboard Interrupt") continue def dispatch_line(self, line: str, prog_name: str = None): diff --git a/pwncat/commands/download.py b/pwncat/commands/download.py index 63dea03..d0ac2e1 100644 --- a/pwncat/commands/download.py +++ b/pwncat/commands/download.py @@ -11,11 +11,22 @@ from pwncat.commands.base import ( from functools import partial from colorama import Fore from pwncat import util +from pwncat.util import console import argparse import datetime import time import os +from rich.progress import ( + BarColumn, + DownloadColumn, + TextColumn, + TransferSpeedColumn, + TimeRemainingColumn, + Progress, + TaskID, +) + class Command(CommandDefinition): """ Download a file from the remote host to the local host""" @@ -23,30 +34,51 @@ class Command(CommandDefinition): PROG = "download" ARGS = { "source": Parameter(Complete.REMOTE_FILE), - "destination": Parameter(Complete.LOCAL_FILE), + "destination": Parameter(Complete.LOCAL_FILE, nargs="?"), } def run(self, args): + # Create a progress bar for the download + progress = Progress( + TextColumn("[bold cyan]{task.fields[filename]}", justify="right"), + BarColumn(bar_width=None), + "[progress.percentage]{task.percentage:>3.1f}%", + "•", + DownloadColumn(), + "•", + TransferSpeedColumn(), + "•", + TimeRemainingColumn(), + ) + + if not args.destination: + args.destination = os.path.basename(args.source) + elif os.path.isdir(args.destination): + args.destination = os.path.join( + args.destination, os.path.basename(args.source) + ) + try: length = pwncat.victim.get_file_size(args.source) started = time.time() - with open(args.destination, "wb") as destination: - with pwncat.victim.open(args.source, "rb", length=length) as source: - util.with_progress( - [ - ("", "downloading "), - ("fg:ansigreen", args.source), - ("", " to "), - ("fg:ansired", args.destination), - ], - partial(util.copyfileobj, source, destination), - length=length, - ) - elapsed = time.time() - started - util.success( - f"downloaded {Fore.CYAN}{util.human_readable_size(length)}{Fore.RESET} " - f"in {Fore.GREEN}{util.human_readable_delta(elapsed)}{Fore.RESET}" + with progress: + task_id = progress.add_task( + "download", filename=args.source, total=length, start=False + ) + with open(args.destination, "wb") as destination: + with pwncat.victim.open(args.source, "rb", length=length) as source: + progress.start_task(task_id) + util.copyfileobj( + source, + destination, + lambda count: progress.update(task_id, advance=count), + ) + elapsed = time.time() - started + + console.log( + f"downloaded [cyan]{util.human_readable_size(length)}[/cyan] " + f"in [green]{util.human_readable_delta(elapsed)}[/green]" ) except (FileNotFoundError, PermissionError, IsADirectoryError) as exc: self.parser.error(str(exc)) diff --git a/pwncat/commands/upload.py b/pwncat/commands/upload.py index dbdef45..cdb6df2 100644 --- a/pwncat/commands/upload.py +++ b/pwncat/commands/upload.py @@ -4,9 +4,24 @@ import time from functools import partial from colorama import Fore +from rich.progress import ( + BarColumn, + DownloadColumn, + TextColumn, + TransferSpeedColumn, + TimeRemainingColumn, + Progress, + TaskID, +) import pwncat -from pwncat import util +from pwncat.util import ( + console, + Access, + human_readable_size, + human_readable_delta, + copyfileobj, +) from pwncat.commands.base import ( CommandDefinition, Complete, @@ -21,35 +36,59 @@ class Command(CommandDefinition): PROG = "upload" ARGS = { "source": Parameter(Complete.LOCAL_FILE), - "destination": Parameter( - Complete.REMOTE_FILE, - type=("method", RemoteFileType(file_exist=False, directory_exist=True)), - ), + "destination": Parameter(Complete.REMOTE_FILE, nargs="?",), } def run(self, args): + # Create a progress bar for the download + progress = Progress( + TextColumn("[bold cyan]{task.fields[filename]}", justify="right"), + BarColumn(bar_width=None), + "[progress.percentage]{task.percentage:>3.1f}%", + "•", + DownloadColumn(), + "•", + TransferSpeedColumn(), + "•", + TimeRemainingColumn(), + ) + + if not args.destination: + args.destination = f"./{os.path.basename(args.source)}" + else: + access = pwncat.victim.access(args.destination) + if Access.DIRECTORY in access: + args.destination = os.path.join( + args.destination, os.path.basename(args.source) + ) + elif Access.PARENT_EXIST not in access: + console.log( + f"[cyan]{args.destination}[/cyan]: no such file or directory" + ) + return + try: length = os.path.getsize(args.source) started = time.time() - with open(args.source, "rb") as source: - with pwncat.victim.open( - args.destination, "wb", length=length - ) as destination: - util.with_progress( - [ - ("", "uploading "), - ("fg:ansigreen", args.source), - ("", " to "), - ("fg:ansired", args.destination), - ], - partial(util.copyfileobj, source, destination), - length=length, - ) + with progress: + task_id = progress.add_task( + "upload", filename=args.destination, total=length, start=False + ) + with open(args.source, "rb") as source: + with pwncat.victim.open( + args.destination, "wb", length=length + ) as destination: + progress.start_task(task_id) + copyfileobj( + source, + destination, + lambda count: progress.update(task_id, advance=count), + ) elapsed = time.time() - started - util.success( - f"uploaded {Fore.CYAN}{util.human_readable_size(length)}{Fore.RESET} " - f"in {Fore.GREEN}{util.human_readable_delta(elapsed)}{Fore.RESET}" + console.log( + f"uploaded [cyan]{human_readable_size(length)}[/cyan] " + f"in [green]{human_readable_delta(elapsed)}[/green]" ) except (FileNotFoundError, PermissionError, IsADirectoryError) as exc: self.parser.error(str(exc)) diff --git a/pwncat/data/gtfobins.json b/pwncat/data/gtfobins.json index aeb1aef..998eedc 100644 --- a/pwncat/data/gtfobins.json +++ b/pwncat/data/gtfobins.json @@ -36,7 +36,8 @@ // to the remote process should be in base64 form, and the // tty is not set to raw mode. // - hex -> same as base64, but base16 instead. - "stream": "raw" + "stream": "raw", + "exit": "{ctrl_c}" }, { "type": "write", diff --git a/pwncat/file.py b/pwncat/file.py index b53f98d..61c8ac9 100644 --- a/pwncat/file.py +++ b/pwncat/file.py @@ -47,6 +47,9 @@ class RemoteBinaryPipe(RawIOBase): if self.exit_cmd and len(self.exit_cmd): pwncat.victim.client.send(self.exit_cmd) + # Flush anything in the queue + pwncat.victim.flush_output() + # Reset the terminal pwncat.victim.restore_remote() # pwncat.victim.reset() @@ -103,17 +106,17 @@ class RemoteBinaryPipe(RawIOBase): piece = self.delim[:i] # if bytes(b[-i:]) == piece: if obj[-i:] == piece: - # try: - # # Peak the next bytes, to see if this is actually the - # # delimeter - # rest = pwncat.victim.client.recv( - # len(self.delim) - len(piece), - # # socket.MSG_PEEK | socket.MSG_DONTWAIT, - # socket.MSG_PEEK, - # ) - # except (socket.error, BlockingIOError): - # rest = b"" - rest = pwncat.victim.peek_output(some=True) + try: + # Peak the next bytes, to see if this is actually the + # delimeter + rest = pwncat.victim.client.recv( + len(self.delim) - len(piece), + # socket.MSG_PEEK | socket.MSG_DONTWAIT, + socket.MSG_PEEK, + ) + except (socket.error, BlockingIOError): + rest = b"" + # rest = pwncat.victim.peek_output(some=True) # It is! if (piece + rest) == self.delim: # Receive the delimeter diff --git a/pwncat/remote/victim.py b/pwncat/remote/victim.py index c424267..be093fc 100644 --- a/pwncat/remote/victim.py +++ b/pwncat/remote/victim.py @@ -1743,7 +1743,7 @@ class Victim: :param some: if true, wait for at least one byte of data before flushing. :type some: bool """ - output = b"" + output = 0 old_timeout = self.client.gettimeout() self.client.settimeout(0) # self.client.send(b"echo\n") @@ -1755,9 +1755,9 @@ class Victim: if len(new) == 0: if len(output) > 0 or some is False: break - output += new + output += len(new) except (socket.timeout, BlockingIOError): - if len(output) > 0 or some is False: + if output > 0 or some is False: break self.client.settimeout(old_timeout) diff --git a/pwncat/util.py b/pwncat/util.py index 6256442..f6d1101 100644 --- a/pwncat/util.py +++ b/pwncat/util.py @@ -21,6 +21,10 @@ import tty import sys import os +from rich.console import Console + +console = Console() + CTRL_C = b"\x03" ALPHANUMERIC = string.ascii_letters + string.digits diff --git a/requirements.txt b/requirements.txt index 0b72ebc..3a063d0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ commentjson requests sqlalchemy pytablewriter +rich diff --git a/setup.py b/setup.py index d8df257..577964a 100644 --- a/setup.py +++ b/setup.py @@ -18,6 +18,7 @@ dependencies = [ "sqlalchemy", "paramiko", "pytablewriter", + "rich", ] dependency_links = [