Added flag

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cloudflare-gateway-adblocking"
-version = "0.1.2"
+version = "0.1.3"
 description = "Serverless adblocking via Cloudflare Zero Trust Gateway"
 authors = ["slastechno <77907286+slashtechno@users.noreply.github.com>"]
 license = "MIT"

src/__main__.py

@@ -35,7 +35,12 @@ def main():
     argparser.add_argument(
         "--log-level", "-l", help="Log level", default="INFO"
     )  # noqa E501
-
+    argparser.add_argument(
+            "--timeout",
+            help="Timeout for requests",
+            default=None,
+            type = int
+            )
     # Credential arguments
     credential_args.add_argument(
         "--account-id",
@@ -121,19 +126,19 @@ def upload_to_cloudflare(args):
     blocklists = upload.get_blocklists(args.blocklists)
     blocklists = upload.apply_whitelists(blocklists, args.whitelists)
     lists = upload.split_list(blocklists)
-    asyncio.run(upload.upload_to_cloudflare(lists, ACCOUNT_ID, TOKEN))
-    cloud_lists = utils.get_lists(ACCOUNT_ID, TOKEN)
+    asyncio.run(upload.upload_to_cloudflare(lists, ACCOUNT_ID, TOKEN, args.timeout))
+    cloud_lists = utils.get_lists(ACCOUNT_ID, TOKEN, args.timeout)
     cloud_lists = utils.filter_adblock_lists(cloud_lists)
-    upload.create_dns_policy(cloud_lists, ACCOUNT_ID, TOKEN)
+    upload.create_dns_policy(cloud_lists, ACCOUNT_ID, TOKEN, args.timeout)
 
 
 def delete_from_cloudflare(args):
     logger.info("Deleting from Cloudflare")
-    rules = utils.get_gateway_rules(ACCOUNT_ID, TOKEN)
-    delete.delete_adblock_policy(rules, ACCOUNT_ID, TOKEN)
-    lists = utils.get_lists(ACCOUNT_ID, TOKEN)
+    rules = utils.get_gateway_rules(ACCOUNT_ID, TOKEN, timeout = args.timeout)
+    delete.delete_adblock_policy(rules, ACCOUNT_ID, TOKEN, args.timeout)
+    lists = utils.get_lists(ACCOUNT_ID, TOKEN, args.timeout)
     lists = utils.filter_adblock_lists(lists)
-    asyncio.run(delete.delete_adblock_list(lists, ACCOUNT_ID, TOKEN))
+    asyncio.run(delete.delete_adblock_list(lists, ACCOUNT_ID, TOKEN, args.timeout))
 
 
 def set_primary_logger(log_level):

src/utils/delete.py

@@ -6,7 +6,7 @@ import requests
 from . import utils
 
 
-async def delete_adblock_list(lists: dict, account_id: str, token: str):
+async def delete_adblock_list(lists: dict, account_id: str, token: str, timeout:int = 10):
     try:
         async with httpx.AsyncClient() as client:
             for lst in lists:
@@ -15,7 +15,7 @@ async def delete_adblock_list(lists: dict, account_id: str, token: str):
                     "Authorization": f"Bearer {token}",
                     "Content-Type": "application/json",
                 }
-                response = await client.delete(url, headers=headers, timeout=10)
+                response = await client.delete(url, headers=headers, timeout=timeout)
                 if response.status_code != 200:
                     print(f"Error deleting list: {response.text}")
                 else:
@@ -27,7 +27,7 @@ async def delete_adblock_list(lists: dict, account_id: str, token: str):
             raise e
 
 
-def delete_adblock_policy(policies: dict, account_id: str, token: str):
+def delete_adblock_policy(policies: dict, account_id: str, token: str, timeout:int = 10):
     for policy in policies:
         if policy["name"] == "Block Ads":
             url = f'https://api.cloudflare.com/client/v4/accounts/{account_id}/gateway/rules/{policy["id"]}'
@@ -35,7 +35,7 @@ def delete_adblock_policy(policies: dict, account_id: str, token: str):
                 "Authorization": f"Bearer {token}",
                 "Content-Type": "application/json",
             }
-            response = requests.delete(url, headers=headers, timeout=10)
+            response = requests.delete(url, headers=headers, timeout=timeout)
             if response.status_code != 200:
                 print(f"Error deleting policy: {response.text}")
             else:

src/utils/upload.py

@@ -48,7 +48,7 @@ def split_list(blocklists):
     return lists
 
 
-async def upload_to_cloudflare(lists, account_id: str, token: str) -> None:
+async def upload_to_cloudflare(lists, account_id: str, token: str, timeout:int = 10) -> None:
     async with httpx.AsyncClient() as client:
         for i, lst in enumerate(lists):
             list_name = f"adblock-list-{i + 1}"
@@ -70,14 +70,14 @@ async def upload_to_cloudflare(lists, account_id: str, token: str) -> None:
                     for x in lst
                 ],
             }
-            response = await client.post(url, headers=headers, json=data)
+            response = await client.post(url, headers=headers, json=data, timeout=timeout)
             print(f"Uploaded {list_name} to Cloudflare")
             if response.status_code != 200:
                 print(f"Error uploading {list_name}: {response.text}")
                 exit(1)
 
 
-def create_dns_policy(lists, account_id: str, token: str) -> None:
+def create_dns_policy(lists, account_id: str, token: str, timeout = 10) -> None:
     url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/gateway/rules"
     headers = {
         "Authorization": f"Bearer {token}",
@@ -98,7 +98,7 @@ def create_dns_policy(lists, account_id: str, token: str) -> None:
         "traffic": traffic,
         "enabled": True,
     }
-    response = requests.post(url, headers=headers, json=data, timeout=10)
+    response = requests.post(url, headers=headers, json=data, timeout=timeout)
     if response.status_code != 200:
         print(f"Error creating DNS policy: {response.text}")
 

src/utils/utils.py

@@ -85,13 +85,13 @@ def convert_to_list(file: pathlib.Path) -> list:
 # General Utils
 
 
-def get_lists(account_id, token) -> dict:
+def get_lists(account_id, token, timeout = 10) -> dict:
     url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/gateway/lists"
     headers = {
         "Authorization": f"Bearer {token}",
         "Content-Type": "application/json",
     }
-    response = requests.get(url, headers=headers, timeout=10)
+    response = requests.get(url, headers=headers, timeout=timeout)
     if response.status_code != 200:
         print(f"Error getting lists: {response.text}")
     return response.json()["result"]
@@ -111,13 +111,13 @@ def filter_adblock_lists(lists: dict) -> dict:
     return adblock_lists
 
 
-def get_gateway_rules(account_id, token) -> dict:
+def get_gateway_rules(account_id, token, timeout = 10) -> dict:
     url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/gateway/rules"
     headers = {
         "Authorization": f"Bearer {token}",
         "Content-Type": "application/json",
     }
-    response = requests.get(url, headers=headers, timeout=10)
+    response = requests.get(url, headers=headers, timeout=timeout)
     if response.status_code != 200:
         print(f"Error getting lists: {response.text}")
     return response.json()["result"]