From ce5c516de14fcbf44154845b44d0685c03c1590b Mon Sep 17 00:00:00 2001
From: slashtechno <77907286+slashtechno@users.noreply.github.com>
Date: Mon, 7 Aug 2023 08:25:55 -0400
Subject: [PATCH] Added `--log-level` flag

---
 README.md          |  4 +++-
 src/__main__.py    | 53 ++++++++++++++++++++++++++++------------------
 src/utils/utils.py |  2 +-
 3 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index 6220236..85bfac9 100644
--- a/README.md
+++ b/README.md
@@ -43,4 +43,6 @@ For example:
 #### Deleting blocklists and firewall policy  
 To delete the blocklists from Cloudflare and delete the firewall policy, use the `delete` subcommand.  
 For example:  
-`cloudflare-gateway-adblocking delete`  
\ No newline at end of file
+`cloudflare-gateway-adblocking delete`  
+### Help  
+For help, use the `--help` flag.  
\ No newline at end of file
diff --git a/src/__main__.py b/src/__main__.py
index 5507922..d2660f1 100644
--- a/src/__main__.py
+++ b/src/__main__.py
@@ -16,20 +16,6 @@ ACCOUNT_ID = None
 
 
 def main():
-    # Setup logging
-    logger.remove()
-    # ^10 is a formatting directive to center with a padding of 10
-    logger_format = "<green>{time:YYYY-MM-DD HH:mm:ss}</green> |<level>{level: ^10}</level>| <level>{message}</level>"  # noqa E501
-    logger.add(stderr, format=logger_format, colorize=True, level="DEBUG")
-
-    # Load .env if it exists
-    # This must precede the argparse setup as env variables are used as default values
-    if Path(".env").is_file():
-        dotenv.load_dotenv()
-        logger.info("Loaded .env file")
-    else:
-        logger.info("No .env file found")
-
     # Parse arguments
 
     # Set up argparse
@@ -43,17 +29,20 @@ def main():
     credential_args = argparser.add_argument_group("Cloudflare Credentials")
 
     # Add arguments
+
+    # General arguments
+    argparser.add_argument('--log-level', '-l', help='Log level', default='INFO')  # noqa E501
+
+    # Credential arguments
     credential_args.add_argument(
         "--account-id",
         "-a",
         help="Cloudflare account ID - environment variable: CLOUDFLARE_ACCOUNT_ID",
-        default=os.environ.get("CLOUDFLARE_ACCOUNT_ID"),
     )
     credential_args.add_argument(
         "--token",
         "-t",
         help="Cloudflare API token - environment variable: CLOUDFLARE_TOKEN",
-        default=os.environ.get("CLOUDFLARE_TOKEN"),
     )
 
     # Add subcommands
@@ -88,19 +77,35 @@ def main():
 
     args = argparser.parse_args()
 
+    # Set up logging
+    set_primary_logger(args.log_level)
     logger.debug(args)
 
+
+
     # Load variables
     global TOKEN
     global ACCOUNT_ID
     TOKEN = args.token
     ACCOUNT_ID = args.account_id
-    # Check if variables are set
+
+    # Check if credentials are set, if they are not, attempt to load from environment variables and .env  # noqa E501
     if TOKEN is None or ACCOUNT_ID is None:
-        logger.error(
-            "No environment variables found. Please create a .env file or .envrc file"
-        )  # noqa E501
-        exit(1)
+        logger.info("No credentials provided with flags")
+        if Path(".env").is_file():
+            logger.debug("Loading .env")
+            dotenv.load_dotenv()
+        else:
+            logger.debug("No .env file found")
+        try:
+            TOKEN = os.environ["CLOUDFLARE_TOKEN"]
+            ACCOUNT_ID = os.environ["CLOUDFLARE_ACCOUNT_ID"]
+            logger.info("Loaded credentials from environment variables")
+        except KeyError:
+            logger.error("No credentials provided")
+            argparser.print_help()
+            exit(1)
+
     try:
         args.func(args)
     except AttributeError:
@@ -128,6 +133,12 @@ def delete_from_cloudflare(args):
     lists = utils.filter_adblock_lists(lists)
     delete.delete_adblock_list(lists, ACCOUNT_ID, TOKEN)
 
+def set_primary_logger(log_level):
+    logger.remove()
+    # ^10 is a formatting directive to center with a padding of 10
+    logger_format = "<green>{time:YYYY-MM-DD HH:mm:ss}</green> |<level>{level: ^10}</level>| <level>{message}</level>"  # noqa E501
+    logger.add(stderr, format=logger_format, colorize=True, level=log_level)
+
 
 if __name__ == "__main__":
     main()
diff --git a/src/utils/utils.py b/src/utils/utils.py
index b3eb8fc..99e5d66 100644
--- a/src/utils/utils.py
+++ b/src/utils/utils.py
@@ -78,7 +78,7 @@ def convert_to_list(file: pathlib.Path) -> list:
             for match in matches
             if match and match.group(1) not in loopback
         ]
-        print(f"First 5 hosts: {hosts[:5]}")
+        # print(f"First 5 hosts: {hosts[:5]}")
         return hosts