From a6609c197202ff7dfd0f16833c60cc5c701d4472 Mon Sep 17 00:00:00 2001
From: slashtechno <77907286+slashtechno@users.noreply.github.com>
Date: Tue, 1 Aug 2023 00:06:22 -0500
Subject: [PATCH] Updated `cmd.py` to call functions

---
 .envrc.example => .env.example |   0
 .envrc.bat                     |   4 --
 README.md                      |   2 +
 poetry.lock                    |  16 +++++-
 pyproject.toml                 |   5 +-
 src/cmd.py                     | 102 +++++++++++++++++++++++++++++----
 src/utils/__init.py__          |   0
 src/utils/__init__.py          |   3 +
 src/utils/adblock_zerotrust.py |  11 ++--
 9 files changed, 122 insertions(+), 21 deletions(-)
 rename .envrc.example => .env.example (100%)
 delete mode 100644 .envrc.bat
 create mode 100644 README.md
 delete mode 100644 src/utils/__init.py__
 create mode 100644 src/utils/__init__.py

diff --git a/.envrc.example b/.env.example
similarity index 100%
rename from .envrc.example
rename to .env.example
diff --git a/.envrc.bat b/.envrc.bat
deleted file mode 100644
index 9abd936..0000000
--- a/.envrc.bat
+++ /dev/null
@@ -1,4 +0,0 @@
-@REM This script is used to load environment variables from .envrc file
-@echo off
-for /f "delims=" %%a in ('type .envrc') do set %%a
-echo %ENV_VAR
\ No newline at end of file
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..c9961f7
--- /dev/null
+++ b/README.md
@@ -0,0 +1,2 @@
+# Cloudflare Gateway Adblocking  
+Serverless adblocking via Cloudflare Zero Trust Gateway  
\ No newline at end of file
diff --git a/poetry.lock b/poetry.lock
index 28c0a9a..a1b7fdd 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -135,6 +135,20 @@ win32-setctime = {version = ">=1.0.0", markers = "sys_platform == \"win32\""}
 [package.extras]
 dev = ["Sphinx (==5.3.0)", "colorama (==0.4.5)", "colorama (==0.4.6)", "freezegun (==1.1.0)", "freezegun (==1.2.2)", "mypy (==v0.910)", "mypy (==v0.971)", "mypy (==v0.990)", "pre-commit (==3.2.1)", "pytest (==6.1.2)", "pytest (==7.2.1)", "pytest-cov (==2.12.1)", "pytest-cov (==4.0.0)", "pytest-mypy-plugins (==1.10.1)", "pytest-mypy-plugins (==1.9.3)", "sphinx-autobuild (==2021.3.14)", "sphinx-rtd-theme (==1.2.0)", "tox (==3.27.1)", "tox (==4.4.6)"]
 
+[[package]]
+name = "python-dotenv"
+version = "1.0.0"
+description = "Read key-value pairs from a .env file and set them as environment variables"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "python-dotenv-1.0.0.tar.gz", hash = "sha256:a8df96034aae6d2d50a4ebe8216326c61c3eb64836776504fcca410e5937a3ba"},
+    {file = "python_dotenv-1.0.0-py3-none-any.whl", hash = "sha256:f5971a9226b701070a4bf2c38c89e5a3f0d64de8debda981d1db98583009122a"},
+]
+
+[package.extras]
+cli = ["click (>=5.0)"]
+
 [[package]]
 name = "requests"
 version = "2.31.0"
@@ -216,4 +230,4 @@ dev = ["black (>=19.3b0)", "pytest (>=4.6.2)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "2213d241196503176caefedd103dfc06ddd3a12c386e2948706dc0d732d7ed50"
+content-hash = "843328e3217d5ebc1fbb55737ca6edc4d47894e2b1c48e5ab5f3fe0b6a27045c"
diff --git a/pyproject.toml b/pyproject.toml
index 3e51a94..fb7090f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,15 +2,16 @@
 name = "cloudflare-gateway-adblocking"
 version = "0.1.0"
 description = "Serverless adblocking via Cloudflare Zero Trust Gateway"
-authors = ["slastechno 77907286+slashtechno@users.noreply.github.com"]
+authors = ["slastechno <77907286+slashtechno@users.noreply.github.com>"]
 license = "MIT"
 readme = "README.md"
-packages = [{include = "cloudflare_gateway_adblocking"}]
+packages = [{include = "src"}]
 
 [tool.poetry.dependencies]
 python = "^3.10"
 requests = "^2.31.0"
 loguru = "^0.7.0"
+python-dotenv = "^1.0.0"
 
 
 [tool.poetry.group.dev.dependencies]
diff --git a/src/cmd.py b/src/cmd.py
index eea55c9..9bc75c3 100644
--- a/src/cmd.py
+++ b/src/cmd.py
@@ -1,39 +1,121 @@
+from loguru import logger 
+# Import the utils package
 import utils
+
 import argparse
 import os
-from loguru import logger
+import dotenv
+from sys import exit, stderr
+from pathlib import Path
+
+TOKEN = None
+ACCOUNT_ID = None
 
 def cli():
+    # 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>"  # pylint: disable=line-too-long
+    logger.add(stderr, format=logger_format, colorize=True, level="DEBUG")
+
+
+    # Load .env if it exists
+    # This must precede the argparse setup as os.environ is used in the default arguments
+    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
-    argparser = argparse.ArgumentParser()
+    argparser = argparse.ArgumentParser(
+        prog='Cloudflare Gateway Adblocking',
+        description='Serverless adblocking via Cloudflare Zero Trust Gateway',
+        epilog=':)'
+    )
+
+    # Add argument groups
+    credential_args = argparser.add_argument_group('Cloudflare Credentials')
+
     # Add arguments
-    argparser.add_argument(
+    credential_args.add_argument(
         "--account-id",
         "-a",
         help="Cloudflare account ID - environment variable: CLOUDFLARE_ACCOUNT_ID",
         default=os.environ.get("CLOUDFLARE_ACCOUNT_ID")
     )
-    argparser.add_argument('--token',
+    credential_args.add_argument('--token',
         '-t', 
         help='Cloudflare API token - environment variable: CLOUDFLARE_TOKEN',
         default=os.environ.get("CLOUDFLARE_TOKEN")
     )
-    # Parse arguments
+
+    # Add subcommands
+    subparsers = argparser.add_subparsers(
+        title="subcommands",
+        description="",
+        help="Subcommands to preform operations",
+        dest="subcommand"
+    )
+    # Add subcommand: upload
+    upload_parser = subparsers.add_parser(
+        "upload",
+        help="Upload adblock lists to Cloudflare"
+    )
+    upload_parser.set_defaults(func=upload_to_cloudflare)
+    upload_parser.add_argument(
+        "--blocklists",
+        "-b",
+        help="Either a blocklist hosts file or a directory containing blocklist hosts files",
+        default="blocklists"
+    )
+    upload_parser.add_argument(
+        "--whitelists",
+        "-w",
+        # help="Either a whitelist hosts file or a directory containing whitelist hosts files"
+        help="Whitelist hosts file or directory",
+        default="whitelist.txt" # Need to change this so it's optional
+    )
+    # Add subcommand: delete
+    delete_parser = subparsers.add_parser(
+        "delete",
+        help="Delete adblock lists from Cloudflare"
+    )
+    delete_parser.set_defaults(func=delete_from_cloudflare)
+
     args = argparser.parse_args()
+
     logger.debug(args)
 
     # Load variables
+    global TOKEN
+    global ACCOUNT_ID
     TOKEN = args.token
     ACCOUNT_ID = args.account_id
-    
     # Check if variables are set
     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()
-    
-
-
+        exit(1)
+    args.func(args)
 
+def upload_to_cloudflare(args):
+    logger.info("Uploading to Cloudflare")
+    blocklists = utils.utils.get_blocklists(args.blocklists)
+    blocklists = utils.adblock_zerotrust.apply_whitelists(blocklists, args.whitelists)
+    lists = utils.adblock_zerotrust.split_list(blocklists)
+    utils.adblock_zerotrust.upload_to_cloudflare(lists, ACCOUNT_ID, TOKEN)
+    cloud_lists = utils.utils.get_lists(ACCOUNT_ID, TOKEN)
+    cloud_lists = utils.utils.filter_adblock_lists(cloud_lists)
+    utils.adblock_zerotrust.create_dns_policy(cloud_lists, ACCOUNT_ID, TOKEN)
+def delete_from_cloudflare(args):
+    logger.info("Deleting from Cloudflare")
+    rules = utils.utils.get_gateway_rules(ACCOUNT_ID, TOKEN)
+    utils.delete_adblock_zerotrust.delete_adblock_policy(rules, ACCOUNT_ID, TOKEN)
+    lists = utils.utils.get_lists(ACCOUNT_ID, TOKEN)
+    lists = utils.utils.filter_adblock_lists(lists)
+    utils.delete_adblock_zerotrust.delete_adblock_list(lists, ACCOUNT_ID, TOKEN)
 
 if __name__ == "__main__":
     cli()
diff --git a/src/utils/__init.py__ b/src/utils/__init.py__
deleted file mode 100644
index e69de29..0000000
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
new file mode 100644
index 0000000..b045e2a
--- /dev/null
+++ b/src/utils/__init__.py
@@ -0,0 +1,3 @@
+from .utils import *
+from .adblock_zerotrust import *
+from .delete_adblock_zerotrust import *
diff --git a/src/utils/adblock_zerotrust.py b/src/utils/adblock_zerotrust.py
index 8c2ac3d..6b326f3 100644
--- a/src/utils/adblock_zerotrust.py
+++ b/src/utils/adblock_zerotrust.py
@@ -3,17 +3,20 @@ import utils
 import pathlib
 
 
-def get_blocklists():
+def get_blocklists(path_to_blocklists: str = None):
     # __file__ is a special variable that is the path to the current file
-    list_directory = pathlib.Path(__file__).parent.parent.joinpath("blocklists")
+    # list_directory = pathlib.Path(__file__).parent.parent.joinpath("blocklists")
+
+    list_directory = pathlib.Path(path_to_blocklists)
     for file in list_directory.iterdir():
         blocklists = utils.convert_to_list(file)
     return blocklists
 
 
-def apply_whitelists(blocklists):
+def apply_whitelists(blocklists, whitelist: str = None):
     whitelist = utils.convert_to_list(
-        pathlib.Path(__file__).parent.parent.joinpath("whitelist.txt")
+        # pathlib.Path(__file__).parent.parent.joinpath("whitelist.txt")
+        pathlib.Path(whitelist)
     )
     blocklists = [x for x in blocklists if x not in whitelist]
     return blocklists