From 9febdde025c0c1139a134b0aaba84a959205b4a6 Mon Sep 17 00:00:00 2001
From: Darryl Nixon <git@nixon.mozmail.com>
Date: Wed, 7 Jun 2023 15:40:02 -0700
Subject: [PATCH] Add Rocketry for analytics processing, add uvloop

---
 crowdtls/analytics/main.py   | 11 ++++++++++
 crowdtls/{cli.py => api.py}  |  8 +++++---
 crowdtls/helpers.py          |  6 +++---
 crowdtls/main.py             | 39 ++++++++++++++++++++++++++++++++++++
 crowdtls/scheduler.py        |  6 ++++++
 crowdtls/{api => }/v1/api.py |  0
 pyproject.toml               |  7 ++++++-
 7 files changed, 70 insertions(+), 7 deletions(-)
 create mode 100644 crowdtls/analytics/main.py
 rename crowdtls/{cli.py => api.py} (86%)
 create mode 100644 crowdtls/main.py
 create mode 100644 crowdtls/scheduler.py
 rename crowdtls/{api => }/v1/api.py (100%)

diff --git a/crowdtls/analytics/main.py b/crowdtls/analytics/main.py
new file mode 100644
index 0000000..85aa802
--- /dev/null
+++ b/crowdtls/analytics/main.py
@@ -0,0 +1,11 @@
+from fastapi import Depends
+from rocketry.conds import hourly
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from crowdtls.db import get_session
+from crowdtls.scheduler import app as schedule
+
+
+@schedule.task(hourly)
+async def find_anomalies(session: AsyncSession = Depends(get_session)):
+    pass
diff --git a/crowdtls/cli.py b/crowdtls/api.py
similarity index 86%
rename from crowdtls/cli.py
rename to crowdtls/api.py
index 6816e13..b15109a 100644
--- a/crowdtls/cli.py
+++ b/crowdtls/api.py
@@ -1,14 +1,12 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
-from crowdtls.api.v1.api import app as api_v1_app
 from crowdtls.db import create_db_and_tables
 from crowdtls.logs import logger
+from crowdtls.v1.api import app as api_v1_app
 
 app = FastAPI()
-
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["POST"], allow_headers=["*"])
-
 app.include_router(api_v1_app, prefix="/api/v1")
 
 
@@ -20,3 +18,7 @@ async def startup_event():
     except Exception:
         logger.error("Failed to create database and tables")
         raise
+
+
+if __name__ == "__main__":
+    app.run()
diff --git a/crowdtls/helpers.py b/crowdtls/helpers.py
index 02e6ce0..57a73db 100644
--- a/crowdtls/helpers.py
+++ b/crowdtls/helpers.py
@@ -1,10 +1,10 @@
 from typing import List
 
+import tldextract
 from cryptography import x509
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
 from fastapi import HTTPException
-from tld import get_tld
 
 from crowdtls.logs import logger
 from crowdtls.models import Certificate
@@ -35,8 +35,8 @@ def decode_der(fingerprint: str, raw_der_certificate: List[int]) -> Certificate:
 
 def parse_hostname(hostname: str) -> Domain:
     try:
-        parsed_domain = get_tld(f"https://{hostname}", as_object=True)
-        return Domain(fqdn=hostname, root=parsed_domain.domain, tld=parsed_domain.tld)
+        parsed_domain = tldextract.extract(hostname)
+        return Domain(fqdn=hostname, root=parsed_domain.domain, tld=parsed_domain.suffix)
     except Exception:
         logger.error(f"Failed to parse hostname: {hostname}")
 
diff --git a/crowdtls/main.py b/crowdtls/main.py
new file mode 100644
index 0000000..853ef37
--- /dev/null
+++ b/crowdtls/main.py
@@ -0,0 +1,39 @@
+import asyncio
+import sys
+from types import FrameType
+
+import uvicorn
+import uvloop
+
+from crowdtls.api import app as app_fastapi
+from crowdtls.logs import logger
+from crowdtls.scheduler import app as app_rocketry
+
+
+class CrowdTLS(uvicorn.Server):
+    def handle_exit(self, sig: int, frame: FrameType) -> None:
+        logger.info("Shutting down CrowdTLS")
+        return super().handle_exit(sig, frame)
+
+
+async def start_server():
+    logger.info("Starting CrowdTLS")
+    server = CrowdTLS(config=uvicorn.Config(app=app_fastapi, workers=1, loop="uvloop"))
+
+    fastapi = asyncio.create_task(server.serve())
+    rocket = asyncio.create_task(app_rocketry.serve())
+
+    await asyncio.wait([rocket, fastapi], return_when=asyncio.FIRST_COMPLETED)
+
+
+def run():
+    if sys.version_info >= (3, 11):
+        with asyncio.Runner(loop_factory=uvloop.new_event_loop) as runner:
+            runner.run(start_server())
+    else:
+        uvloop.install()
+        asyncio.run(start_server())
+
+
+if __name__ == "__main__":
+    run()
diff --git a/crowdtls/scheduler.py b/crowdtls/scheduler.py
new file mode 100644
index 0000000..fdf4c12
--- /dev/null
+++ b/crowdtls/scheduler.py
@@ -0,0 +1,6 @@
+from rocketry import Rocketry
+
+app = Rocketry(execution="async")
+
+if __name__ == "__main__":
+    app.run()
diff --git a/crowdtls/api/v1/api.py b/crowdtls/v1/api.py
similarity index 100%
rename from crowdtls/api/v1/api.py
rename to crowdtls/v1/api.py
diff --git a/pyproject.toml b/pyproject.toml
index 861a277..b4aa11d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,9 +19,14 @@ dependencies = [
     "greenlet==2.0.2",
     "sqlmodel==0.0.8",
     "sqlalchemy==1.4.41",
-    "tld>=0.13",
+    "tldextract>=3.4.4",
+    "rocketry>=2.5.1",
+    "uvloop>=0.17.0",
 ]
 
+[project.scripts]
+crowdtls = "crowdtls:main.run"
+
 [project.urls]
 homepage = "https://github.com/DarrylNixon/CrowdTLS"
 repository = "https://github.com/DarrylNixon/CrowdTLS-server"