Merge pull request #207 from bee-san/improve-printing

pywhat/printer.py

@@ -1,5 +1,6 @@
 import json
 import os
+import re
 
 from rich.console import Console
 from rich.table import Table
@@ -10,7 +11,7 @@ def __init__(self):
         self.console = Console(highlight=False)
         self.bug_bounty_mode = False
 
-    def pretty_print(self, text: dict, text_input):
+    def pretty_print(self, text: dict, text_input, print_tags=False):
         to_out = ""
 
         if text["File Signatures"]:
@@ -69,7 +70,14 @@ def pretty_print(self, text: dict, text_input):
                     ):
                         exploit = i["Regex Pattern"]["Exploit"]
 
-                    if not description:
+                    if print_tags:
+                        tags = f"Tags: {', '.join(i['Regex Pattern']['Tags'])}"
+                        if description is None:
+                            description = tags
+                        else:
+                            description += "\n" + tags
+
+                    if description is None:
                         description = "None"
 
                     # FIXME this is quite messy
@@ -116,7 +124,7 @@ def print_json(self, text: dict):
     Returns the printable object
     """
 
-    def print_raw(self, text: dict, text_input) -> str:
+    def print_raw(self, text: dict, text_input, print_tags=False):
         output_str = ""
 
         if text["File Signatures"] and text["Regexes"]:
@@ -170,6 +178,10 @@ def print_raw(self, text: dict, text_input) -> str:
                             "\n[bold #D7Afff]Exploit: [/bold #D7Afff]"
                             + i["Regex Pattern"]["Exploit"]
                         )
+
+                    if print_tags:
+                        output_str += f"\n[bold #D7Afff]Tags: [/bold #D7Afff]{', '.join(i['Regex Pattern']['Tags'])}"
+
                     output_str += "\n\n"
 
         if output_str == "" and not self.bug_bounty_mode:
@@ -180,6 +192,48 @@ def print_raw(self, text: dict, text_input) -> str:
 
         return output_str
 
+    def format_print(self, text: dict, format_str: str):
+        if text["Regexes"]:
+            output = []
+            format_list = []
+
+            # Split format_str so that format_list's item will either be r'\\' or something else
+            start = 0
+            i = format_str.find(r"\\", start)
+            while i != -1:
+                if format_str[start:i]:
+                    format_list.append(format_str[start:i])
+                format_list.append("\\")
+                start = i + 2
+                i = format_str.find(r"\\", start)
+            format_list.append(format_str[start:])
+
+            for key, value in text["Regexes"].items():
+                for match in value:
+                    temp = ""
+                    for s in format_list:
+                        formats = {
+                            "%m": match["Matched"],
+                            "%n": match["Regex Pattern"]["Name"],
+                            "%d": match["Regex Pattern"]["Description"],
+                            "%e": match["Regex Pattern"].get("Exploit"),
+                            "%r": str(match["Regex Pattern"]["Rarity"]),
+                            "%l": match["Regex Pattern"]["URL"] + match["Matched"]
+                            if match["Regex Pattern"]["URL"] is not None
+                            else None,
+                            "%t": ", ".join(match["Regex Pattern"]["Tags"]),
+                        }
+                        for format, value in formats.items():
+                            value = str() if value is None else value
+                            s = re.sub(r"(?<!\\)" + format, value, s)
+                        s = re.sub(r"\\%", "%", s)
+                        temp += s
+                    output.append(temp)
+
+            str_output = "\n".join(output)
+            if str_output.strip():
+                self.console.print(str_output)
+
     def _check_if_exploit_in_json(self, text: dict) -> bool:
         if "File Signatures" in text and text["File Signatures"]:
             # loops files

pywhat/what.py

@@ -118,8 +118,9 @@ def get_text(ctx, opts, value):
 @click.option(
     "--format",
     required=False,
-    help="--format json for json output. --format pretty for a pretty table output.",
+    help="Format output according to specified rules.",
 )
+@click.option("-pt", "--print-tags", is_flag=True, help="Add flags to ouput")
 def main(**kwargs):
     """
     pyWhat - Identify what something is.
@@ -184,6 +185,44 @@ def main(**kwargs):
 
         Refer to the Filtration section for more information.
 
+    Formatting the output:
+
+        --format format_str
+
+            format_str can be equal to:
+
+                pretty - Output data in the table
+
+                json - Ouput data in json format
+
+                CUSTOM_STRING - Print data in the way you want. For every match CUSTOM_STRING will be printed and '%x' (See below for possible x values) will be substituted with a match value.
+
+                For example:
+
+                    pywhat --format '%m - %n' 'google.com htb{flag}'
+
+                    will print:
+
+                    htb{flag} - HackTheBox Flag Format
+                    google.com - Uniform Resource Locator (URL)
+
+                Possible '%x' values:
+
+                    %m - matched text
+
+                    %n - name of regex
+
+                    %d - description (will not output if absent)
+
+                    %e - exploit (will not ouput if absent)
+
+                    %r - rarity
+
+                    %l - link (will not ouput if absent)
+
+                    %t - tags (in 'tag1, tag2 ...' format)
+
+                If you want to print '%' or '\' character - escape it: '\%', '\\'.
 
     Examples:
 
@@ -242,12 +281,14 @@ def main(**kwargs):
 
     p = printer.Printing()
 
-    if kwargs["json"] or kwargs["format"] == "json":
+    if kwargs["json"] or str(kwargs["format"]).strip() == "json":
         p.print_json(identified_output)
-    elif kwargs["format"] == "pretty":
-        p.pretty_print(identified_output, kwargs["text_input"])
+    elif str(kwargs["format"]).strip() == "pretty":
+        p.pretty_print(identified_output, kwargs["text_input"], kwargs["print_tags"])
+    elif kwargs["format"] is not None:
+        p.format_print(identified_output, kwargs["format"])
     else:
-        p.print_raw(identified_output, kwargs["text_input"])
+        p.print_raw(identified_output, kwargs["text_input"], kwargs["print_tags"])
 
 
 class What_Object: