Clean Coding: Text Coloring

gato/attack/attack.py

@@ -77,7 +77,7 @@ def __setup_user_info(self):
             )
             Output.info(
                 "The GitHub Classic PAT has the following scopes: "
-                f'{Output.yellow(", ".join(self.user_perms["scopes"]))}'
+                f'{Output.text_color(", ".join(self.user_perms["scopes"]), "yellow")}'
             )
 
         return True

gato/cli/output.py

@@ -1,13 +1,10 @@
 import json
-
-from gato.cli import (RED_DASH, GREEN_PLUS, GREEN_EXCLAIM, RED_EXCLAIM,
-                      BRIGHT_DASH, YELLOW_EXCLAIM, SPLASH, YELLOW_DASH)
-
-
 from colorama import Style, Fore
+from gato.cli import (RED_DASH, GREEN_PLUS, GREEN_EXCLAIM, RED_EXCLAIM, BRIGHT_DASH, YELLOW_EXCLAIM, SPLASH, YELLOW_DASH)
 
 
 class Singleton (type):
+
     _instances = {}
 
     def __call__(cls, *args, **kwargs):
@@ -17,20 +14,23 @@ def __call__(cls, *args, **kwargs):
             )
         return cls._instances[cls]
 
-
 class Output(metaclass=Singleton):
 
     def __init__(self, silent: bool, color: bool):
+
         self.silent = silent
         self.color = color
-
         self.red_dash = RED_DASH if color else '[-]'
         self.red_explain = RED_EXCLAIM if color else '[!]'
         self.green_plus = GREEN_PLUS if color else '[+]'
         self.green_exclaim = GREEN_EXCLAIM if color else '[!]'
         self.bright_dash = BRIGHT_DASH if color else '-'
         self.yellow_exclaim = YELLOW_EXCLAIM if color else "[!]"
         self.yellow_dash = YELLOW_DASH if color else "[-]"
+
+    #######################
+    ## FileIO  Functions ##
+    #######################
 
     @classmethod
     def write_json(cls, execution_wrapper, output_json):
@@ -43,17 +43,23 @@ def write_json(cls, execution_wrapper, output_json):
         Returns:
             True if successful, false otherwise.
         """
+
         if execution_wrapper.user_details:
             with open(output_json, 'w') as json_out:
                 json_out.write(
                     json.dumps(execution_wrapper.toJSON(), indent=4)
                 )
             return True
+
+    #######################
+    ## STDOUT Functions  ##
+    #######################
 
     @classmethod
     def splash(cls):
         """Prints the Gato mascot.
         """
+
         if not Output().silent:
             print(SPLASH)
 
@@ -64,6 +70,7 @@ def error(cls, message: str):
         Args:
             message (str): Message to format.
         """
+
         print(f"{Output().red_dash} {message}")
 
     @classmethod
@@ -73,6 +80,7 @@ def info(cls, message: str, end='\n', flush=False):
         Args:
             message (str): The message to print.
         """
+
         print(f"{Output().green_plus} {message}", end=end, flush=flush)
 
     @classmethod
@@ -82,6 +90,7 @@ def tabbed(cls, message: str):
         Args:
             message (str): The message to print.
         """
+
         print(f"    {Output().bright_dash} {message}")
 
     @classmethod
@@ -91,6 +100,7 @@ def header(cls, message: str):
         Args:
             message (str): The message to print.
         """
+
         print(
             f"{cls.bright('---')}"
             f" {message} "
@@ -104,6 +114,7 @@ def result(cls, message: str):
         Args:
             message (str): The message to print.
         """
+
         print(f"{Output().green_plus} {message}")
 
     @classmethod
@@ -114,6 +125,7 @@ def owned(cls, message: str):
         Args:
             message (str): The message to print.
         """
+
         print(f"{Output().green_exclaim} {message}")
 
     @classmethod
@@ -131,8 +143,13 @@ def warn(cls, message: str):
         """Used to let the user know something that they should not, but
         unlikely to lead to an exploit.
         """
+
         print(f"{Output().yellow_exclaim} {message}")
 
+
+    #######################
+    ##  Color Functions  ##
+    #######################
     @classmethod
     def bright(cls, toformat: str):
         """Highlights the text and returns it.
@@ -148,33 +165,22 @@ def bright(cls, toformat: str):
             return f'{Style.BRIGHT}{toformat}{Style.RESET_ALL}'
         else:
             return toformat
-
+        
     @classmethod
-    def yellow(cls, toformat: str):
-        """Makes the text yellow and returns it.
-
-        Args:
-            toformat (str): Message to format.
-
-        Returns:
-            (str)): Formatted message.
-        """
-        if cls not in cls._instances or Output().color:
-            return f'{Fore.YELLOW}{toformat}{Style.RESET_ALL}'
-        else:
-            return toformat
-
-    @classmethod
-    def green(cls, toformat: str):
+    def text_color(cls, toformat: str, color: str):
         """Makes the text green and returns it.
 
         Args:
             toformat (str): Message to format.
+            color (str): Message to color. Either
 
         Returns:
             (str)): Formatted message.
         """
+
         if cls not in cls._instances or Output().color:
-            return f'{Fore.GREEN}{toformat}{Style.RESET_ALL}'
+                color = color.upper()
+                color = "Fore." + color
+                return f'{color}{toformat}{Style.RESET_ALL}'
         else:
-            return toformat
+            return toformat

gato/enumerate/enumerate.py

@@ -75,7 +75,7 @@ def __setup_user_info(self):
             if len(self.user_perms["scopes"]):
                 Output.info(
                     "The GitHub Classic PAT has the following scopes: "
-                    f'{Output.yellow(", ".join(self.user_perms["scopes"]))}'
+                    f'{Output.text_color(", ".join(self.user_perms["scopes"]), "yellow")}'
                 )
             else:
                 Output.warn("The token has no scopes!")

gato/enumerate/recommender.py

@@ -211,7 +211,7 @@ def print_org_findings(scopes, organization: Organization):
             Output.owned("The user is an organization owner!")
             if "admin:org" in scopes:
                 Output.result(
-                    f"The token also has the {Output.yellow('admin:org')} "
+                    f"The token also has the {Output.text_color('admin:org', 'yellow')} "
                     "scope. This token has extensive access to the GitHub"
                     " organization!"
                 )

gato/search/search.py

@@ -52,7 +52,7 @@ def __setup_user_info(self):
             if len(self.user_perms["scopes"]) > 0:
                 Output.info(
                     f"The GitHub Classic PAT has the following scopes: "
-                    f'{Output.yellow(", ".join(self.user_perms["scopes"]))}'
+                    f'{Output.text_color(", ".join(self.user_perms["scopes"]), "yellow")}'
                 )
             else:
                 Output.warn("The token has no scopes!")