Saverr/pwncat
1
0
Fork 0
mirror of https://github.com/calebstewart/pwncat.git synced 2024-12-03 13:54:15 +01:00
Code Projects Releases Wiki Activity

Added the rich module

Browse Source 
rich provides better progress bars and log output and exception tracebacks.
This commit is contained in:
Caleb Stewart 2020-06-05 21:32:24 -04:00
parent cf5d809eda
commit 3678e9fa66
9 changed files with 141 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
pwncat/commands/__init__.py
View File

@ -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):

66
pwncat/commands/download.py
View File

@ -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))

83
pwncat/commands/upload.py
View File

@ -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))

3
pwncat/data/gtfobins.json
View File

@ -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",

25
pwncat/file.py
View File

@ -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

6
pwncat/remote/victim.py
View File

@ -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)

4
pwncat/util.py
View File

@ -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

1
requirements.txt
View File

@ -9,3 +9,4 @@ commentjson
requests
sqlalchemy
pytablewriter
rich

1
setup.py
View File

@ -18,6 +18,7 @@ dependencies = [
"sqlalchemy",
"paramiko",
"pytablewriter",
"rich",
]
dependency_links = [