Update server.py

server.py

@@ -6,7 +6,9 @@
 from DrissionPage import ChromiumPage, ChromiumOptions
 from fastapi import FastAPI, HTTPException, Response
 from pydantic import BaseModel
+from typing import Dict
 
+# Chromium options arguments
 arguments = [
     "-no-first-run",
     "-force-color-profile=srgb",
@@ -25,29 +27,26 @@
 browser_path = "/usr/bin/google-chrome"
 app = FastAPI()
 
-
+# Pydantic model for the response
 class CookieResponse(BaseModel):
-    cookies: dict
-
+    cookies: Dict[str, str]
 
-def isSafeURL(url):
+# Function to check if the URL is safe
+def is_safe_url(url: str) -> bool:
     parsed_url = urlparse(url)
     ip_pattern = re.compile(r"^(127\.0\.0\.1|localhost|0\.0\.0\.0|::1|10\.\d+\.\d+\.\d+|172\.1[6-9]\.\d+\.\d+|172\.2[0-9]\.\d+\.\d+|172\.3[0-1]\.\d+\.\d+|192\.168\.\d+\.\d+)$")
     hostname = parsed_url.hostname
     if (hostname and ip_pattern.match(hostname)) or parsed_url.scheme == "file":
         return False
     return True
 
-
-def bypass_cloudlflare(url, retries):
-    # Set up Chromium options
+# Function to bypass Cloudflare protection
+def bypass_cloudflare(url: str, retries: int) -> ChromiumPage:
     options = ChromiumOptions()
     options.set_paths(browser_path=browser_path).headless(False)
 
-    # Initialize the browser
     driver = ChromiumPage(addr_or_opts=options)
     try:
-        # Bypass
         driver.get(url)
         cf_bypasser = CloudflareBypasser(driver, retries, True)
         cf_bypasser.bypass()
@@ -56,28 +55,27 @@ def bypass_cloudlflare(url, retries):
         driver.quit()
         raise e
 
-
+# Endpoint to get cookies
 @app.get("/cookies", response_model=CookieResponse)
 async def get_cookies(url: str, retries: int = 5):
-    if not isSafeURL(url):
+    if not is_safe_url(url):
         raise HTTPException(status_code=400, detail="Invalid URL")
     try:
-        driver = bypass_cloudlflare(url, retries)
+        driver = bypass_cloudflare(url, retries)
         cookies = driver.cookies(as_dict=True)
         driver.quit()
         return CookieResponse(cookies=cookies)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
+# Endpoint to get HTML content and cookies
 @app.get("/html")
-async def get_cookies(url: str, retries: int = 5):
-    if not isSafeURL(url):
+async def get_html(url: str, retries: int = 5):
+    if not is_safe_url(url):
         raise HTTPException(status_code=400, detail="Invalid URL")
     try:
-        driver = bypass_cloudlflare(url, retries)
+        driver = bypass_cloudflare(url, retries)
         html = driver.html
-
         cookies_json = json.dumps(driver.cookies(as_dict=True))
 
         response = Response(content=html, media_type="text/html")
@@ -87,8 +85,7 @@ async def get_cookies(url: str, retries: int = 5):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
+# Main entry point
 if __name__ == "__main__":
     import uvicorn
-
     uvicorn.run(app, host="0.0.0.0", port=8000)