From 5fd2ce94ab0a66ac2033bd8349727bf029fb72f8 Mon Sep 17 00:00:00 2001 From: slashtechno <77907286+slashtechno@users.noreply.github.com> Date: Sat, 19 Aug 2023 15:05:41 -0400 Subject: [PATCH] Added flag --- pyproject.toml | 2 +- src/__main__.py | 21 +++++++++++++-------- src/utils/delete.py | 8 ++++---- src/utils/upload.py | 8 ++++---- src/utils/utils.py | 8 ++++---- 5 files changed, 26 insertions(+), 21 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7e65fc3..62c0655 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/__main__.py b/src/__main__.py index c23e98b..660ed9b 100644 --- a/src/__main__.py +++ b/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): diff --git a/src/utils/delete.py b/src/utils/delete.py index d830624..9e04e64 100644 --- a/src/utils/delete.py +++ b/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: diff --git a/src/utils/upload.py b/src/utils/upload.py index 8646c25..350b6f7 100644 --- a/src/utils/upload.py +++ b/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}") diff --git a/src/utils/utils.py b/src/utils/utils.py index 39db5cd..a2cba57 100644 --- a/src/utils/utils.py +++ b/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"]