diff --git a/pwncat/__main__.py b/pwncat/__main__.py index 0e03357..854b7d2 100644 --- a/pwncat/__main__.py +++ b/pwncat/__main__.py @@ -9,7 +9,7 @@ from sqlalchemy import exc as sa_exc from sqlalchemy.exc import InvalidRequestError import pwncat -from pwncat import util +from pwncat.util import console from pwncat.remote import Victim @@ -60,9 +60,9 @@ def main(): sys.stdout.flush() except ConnectionResetError: pwncat.victim.restore_local_term() - util.warn("connection reset by remote host") + console.log("[yellow]warning[/yellow]: connection reset by remote host") except SystemExit: - util.success("closing down connection.") + console.log("closing connection") finally: # Restore the shell pwncat.victim.restore_local_term() @@ -71,7 +71,7 @@ def main(): pwncat.victim.session.commit() except InvalidRequestError: pass - util.success("local terminal restored") + console.log("local terminal restored") if __name__ == "__main__": diff --git a/pwncat/commands/persist.py b/pwncat/commands/persist.py index fddbe27..5177f71 100644 --- a/pwncat/commands/persist.py +++ b/pwncat/commands/persist.py @@ -5,7 +5,7 @@ from typing import Dict, Type, Tuple, Iterator from colorama import Fore, Style import pwncat -from pwncat import util +from pwncat.util import console from pwncat.commands.base import CommandDefinition, Complete, Parameter, StoreConstOnce from pwncat.persist import PersistenceMethod, PersistenceError @@ -84,80 +84,65 @@ class Command(CommandDefinition): # List of available persistence methods METHODS: Dict[str, Type["PersistenceMethod"]] = {} - @property - def installed_methods(self) -> Iterator[Tuple[str, str, PersistenceMethod]]: - me = pwncat.victim.current_user - for method in pwncat.victim.persist: - if method.system and method.installed(): - yield (method.name, None, method) - elif not method.system: - if me.id == 0: - for user in pwncat.victim.users: - util.progress(f"checking {method.name} for: {user}") - if method.installed(user): - util.erase_progress() - yield (method.name, user, method) - util.erase_progress() - else: - if method.installed(me.name): - yield (method.name, me.name, method) + def show_status(self): + """ Show the list of installed methods """ + + ninstalled = 0 + for user, method in pwncat.victim.persist.installed: + console.print(f" - {method.format(user)} installed") + ninstalled += 1 + if not ninstalled: + console.log("[yellow]warning[/yellow]: no persistence methods installed") + + def list_methods(self, method): + """ List available methods or help for a specific method """ + + if method: + try: + method = next(pwncat.victim.persist.find(method)) + console.print(f"[underline bold]{method.format()}") + console.print(textwrap.indent(textwrap.dedent(method.__doc__), " ")) + except StopIteration: + console.log(f"[red]error[/red]: {method}: no such persistence method") + else: + for method in pwncat.victim.persist: + console.print(f" - {method.format()}") + + def clean_methods(self): + """ Remove all persistence methods from the victim """ + + util.progress("cleaning persistence methods: ") + for user, method in pwncat.victim.persist.installed: + try: + util.progress(f"cleaning persistance methods: {method.format(user)}") + pwncat.victim.persist.remove(method.name, user) + util.success(f"removed {method.format(user)}") + except PersistenceError as exc: + util.erase_progress() + util.warn( + f"{method.format(user)}: removal failed: {exc}\n", overlay=True + ) + util.erase_progress() def run(self, args): - if args.action == "status": - ninstalled = 0 - for user, method in pwncat.victim.persist.installed: - print(f" - {method.format(user)} installed") - ninstalled += 1 - if not ninstalled: - util.warn( - "no persistence methods observed as " - f"{Fore.GREEN}{pwncat.victim.whoami()}{Fore.RED}" - ) - return - elif args.action == "list": - if args.method: - try: - method = next(pwncat.victim.persist.find(args.method)) - print(f"\033[4m{method.format()}{Style.RESET_ALL}") - print(textwrap.indent(textwrap.dedent(method.__doc__), " ")) - except StopIteration: - util.error(f"{args.method}: no such persistence method") - else: - for method in pwncat.victim.persist: - print(f" - {method.format()}") - return - elif args.action == "clean": - util.progress("cleaning persistence methods: ") - for user, method in pwncat.victim.persist.installed: - try: - util.progress( - f"cleaning persistance methods: {method.format(user)}" - ) - pwncat.victim.persist.remove(method.name, user) - util.success(f"removed {method.format(user)}") - except PersistenceError as exc: - util.erase_progress() - util.warn( - f"{method.format(user)}: removal failed: {exc}\n", overlay=True - ) - util.erase_progress() - return - elif args.method is None: - self.parser.error("no method specified") - return - - # Grab the user we want to install the persistence as - if args.user: - user = args.user - else: - # Default is to install as current user - user = pwncat.victim.whoami() - try: - if args.action == "install": - pwncat.victim.persist.install(args.method, user) + if args.action == "status": + self.show_status() + elif args.action == "list": + self.list_methods(args.method) + elif args.action == "clean": + self.clean_methods() + elif args.action == "install": + pwncat.victim.persist.install( + args.method, args.user if args.user else pwncat.victim.whoami() + ) elif args.action == "remove": - pwncat.victim.persist.remove(args.method, user) + pwncat.victim.persist.remove( + args.method, args.user if args.user else pwncat.victim.whoami() + ) + elif args.method is None: + self.parser.error("no method specified") + return except PersistenceError as exc: - util.error(f"{exc}") + console.log(f"[red]error[/red]: {exc}") diff --git a/pwncat/commands/set.py b/pwncat/commands/set.py index c49dbd6..6ab88cb 100644 --- a/pwncat/commands/set.py +++ b/pwncat/commands/set.py @@ -5,7 +5,7 @@ from sqlalchemy.orm import sessionmaker import pwncat from pwncat.commands.base import CommandDefinition, Complete, Parameter -from pwncat import util +from pwncat.util import console, State class Command(CommandDefinition): @@ -38,17 +38,17 @@ class Command(CommandDefinition): found = False for name, user in pwncat.victim.users.items(): if user.password is not None: - print( - f" - {Fore.GREEN}{user}{Fore.RESET} -> {Fore.RED}{repr(user.password)}{Fore.RESET}" + console.print( + f" - [green]{user}[/green] -> [red]{repr(user.password)}[/red]" ) found = True if not found: - util.warn("no known user passwords") + console.log("[yellow]warning[/yellow]: no known user passwords") else: if args.variable not in pwncat.victim.users: self.parser.error(f"{args.variable}: no such user") - print( - f" - {Fore.GREEN}{args.variable}{Fore.RESET} -> {Fore.RED}{repr(args.value)}{Fore.RESET}" + console.print( + f" - [green]{args.variable}[/green] -> [red]{repr(args.value)}[/red]" ) pwncat.victim.users[args.variable].password = args.value else: @@ -58,9 +58,9 @@ class Command(CommandDefinition): and args.value is not None ): try: - pwncat.victim.state = util.State._member_map_[args.value.upper()] + pwncat.victim.state = State._member_map_[args.value.upper()] except KeyError: - util.error(f"{args.value}: invalid state") + console.log(f"[red]error[/red]: {args.value}: invalid state") elif args.variable is not None and args.value is not None: try: pwncat.victim.config[args.variable] = args.value @@ -79,17 +79,15 @@ class Command(CommandDefinition): ) pwncat.victim.session = pwncat.victim.session_maker() except ValueError as exc: - util.error(str(exc)) + console.log(f"[red]error[/red]: {exc}") elif args.variable is not None: value = pwncat.victim.config[args.variable] - print( - f" {Fore.CYAN}{args.variable}{Fore.RESET} = " - f"{Fore.YELLOW}{repr(value)}{Fore.RESET}" + console.print( + f" [cyan]{args.variable}[/cyan] = [yellow]{repr(value)}[/yellow]" ) else: for name in pwncat.victim.config: value = pwncat.victim.config[name] - print( - f" {Fore.CYAN}{name}{Fore.RESET} = " - f"{Fore.YELLOW}{repr(value)}{Fore.RESET}" + console.print( + f" [cyan]{name}[/cyan] = [yellow]{repr(value)}[/yellow]" ) diff --git a/pwncat/commands/sync.py b/pwncat/commands/sync.py index fd39691..afac8a3 100644 --- a/pwncat/commands/sync.py +++ b/pwncat/commands/sync.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import pwncat from pwncat.commands.base import CommandDefinition, Complete, Parameter -from pwncat import util +from pwncat.util import console import os @@ -25,7 +25,9 @@ class Command(CommandDefinition): TERM = os.environ.get("TERM", None) if TERM is None: if not args.quiet: - util.warn("no local TERM set. falling back to 'xterm'") + console.log( + "[yellow]warning[/yellow]: no local [blue]TERM[/blue]; falling back to 'xterm'" + ) TERM = "xterm" # Get the width and height @@ -33,8 +35,11 @@ class Command(CommandDefinition): # Update the state pwncat.victim.run( - f"stty rows {rows};" f"stty columns {columns};" f"export TERM='{TERM}'" + f"stty rows {rows}; stty columns {columns}; export TERM='{TERM}'" ) if not args.quiet: - util.success("terminal state synchronized") + console.log( + "[green]:heavy_check_mark:[/green] terminal state synchronized", + emoji=True, + ) diff --git a/pwncat/commands/tamper.py b/pwncat/commands/tamper.py index 5ff1c04..1b08ce1 100644 --- a/pwncat/commands/tamper.py +++ b/pwncat/commands/tamper.py @@ -1,6 +1,10 @@ #!/usr/bin/env python3 +from typing import List + +from rich.progress import Progress, BarColumn + import pwncat -from pwncat import util +from pwncat.util import console from pwncat.commands.base import ( CommandDefinition, Complete, @@ -8,7 +12,7 @@ from pwncat.commands.base import ( StoreConstOnce, StoreForAction, ) -from pwncat.tamper import RevertFailed +from pwncat.tamper import RevertFailed, Tamper class Command(CommandDefinition): @@ -49,29 +53,36 @@ class Command(CommandDefinition): if args.action == "revert": if args.all: - removed_tampers = [] - util.progress(f"reverting tamper") - for tamper in pwncat.victim.tamper: - try: - util.progress(f"reverting tamper: {tamper}") - tamper.revert() - removed_tampers.append(tamper) - except RevertFailed as exc: - util.warn(f"{tamper}: revert failed: {exc}") - for tamper in removed_tampers: - pwncat.victim.tamper.remove(tamper) - util.success("tampers reverted!") - pwncat.victim.session.commit() + tampers = list(pwncat.victim.tamper) else: - if args.tamper not in range(len(pwncat.victim.tamper)): - self.parser.error("invalid tamper id") - tamper = pwncat.victim.tamper[args.tamper] try: + tampers = [pwncat.victim.tamper[args.tamper]] + except KeyError: + console.log("[red]error[/red]: invalid tamper id") + return + self.revert(tampers) + else: + for ident, tamper in enumerate(pwncat.victim.tamper): + console.print(f" [cyan]{ident}[/cyan] - {tamper}") + + def revert(self, tampers: List[Tamper]): + """ Revert the list of tampers with a nice progress bar """ + + with Progress( + "[bold]reverting[/bold]", + "•", + "{task.fields[tamper]}", + BarColumn(bar_width=None), + "[progress.percentage]{task.percentage:>3.1f}%", + console=console, + ) as progress: + task = progress.add_task("reverting", tamper="init", total=len(tampers)) + for tamper in tampers: + try: + progress.update(task, tamper=str(tamper)) tamper.revert() pwncat.victim.tamper.remove(tamper) except RevertFailed as exc: - util.error(f"revert failed: {exc}") - pwncat.victim.session.commit() - else: - for id, tamper in enumerate(pwncat.victim.tamper): - print(f" {id} - {tamper}") + progress.log(f"[yellow]warning[/yellow]: revert failed: {exc}") + progress.update(task, advance=1) + progress.update(task, tamper="complete") diff --git a/pwncat/privesc/__init__.py b/pwncat/privesc/__init__.py index 1a2a9e6..cb18026 100644 --- a/pwncat/privesc/__init__.py +++ b/pwncat/privesc/__init__.py @@ -630,14 +630,14 @@ class Finder: privkey = None elif writers: # TODO this needs to be updated to work in the middle of a rich progress - util.warn( - "no readers found for {Fore.GREEN}{techniques[0].user}{Fore.RESET}" + console.log( + f"[yellow]warning[/yellow]: no readers found for [green]{techniques[0].user}[/green] " + f"however, we do have a writer." ) - util.warn(f"however, we do have a writer.") - response = confirm( - "would you like to clobber their authorized keys? ", suffix="(y/N) " - ) - if not response: + response = console.input( + "Would you like to clobber their authorized keys? (y/N) " + ).lower() + if response != "y": raise PrivescError("user aborted key clobbering") # If we don't already know a private key, then we need a writer diff --git a/pwncat/tamper.py b/pwncat/tamper.py index 63d5b7e..8ba4dd2 100644 --- a/pwncat/tamper.py +++ b/pwncat/tamper.py @@ -42,7 +42,7 @@ class CreatedFile(Tamper): raise RevertFailed(str(exc)) def __str__(self): - return f"{Fore.RED}Created{Fore.RESET} file {Fore.CYAN}{self.path}{Fore.RESET}" + return f"[red]Created[/red] file [cyan]{self.path}[/cyan]" class ModifiedFile(Tamper): @@ -90,7 +90,7 @@ class ModifiedFile(Tamper): raise RevertFailed(str(exc)) def __str__(self): - return f"{Fore.RED}Modified{Fore.RESET} {Fore.CYAN}{self.path}{Fore.RESET}" + return f"[red]Modified[/red] [cyan]{self.path}[/cyan]" def __repr__(self): return f"ModifiedFile(path={self.path})"