More changed logging

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__":

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}")

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]"
                     )

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,
+            )

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

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

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})"