From 677cfbc420e0cf2fb6c902db77d9ec6c4deb47cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9ophile=20Diot?= Date: Thu, 14 Mar 2024 16:06:09 +0000 Subject: [PATCH] Refactor BunkerWeb Pro license check and metadata storage --- .../core/pro/jobs/download-pro-plugins.py | 25 +++++++++---------- src/common/core/pro/plugin.json | 2 +- src/common/db/Database.py | 5 +++- src/common/db/model.py | 1 + 4 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/common/core/pro/jobs/download-pro-plugins.py b/src/common/core/pro/jobs/download-pro-plugins.py index f8c657287..c3c09fd4c 100644 --- a/src/common/core/pro/jobs/download-pro-plugins.py +++ b/src/common/core/pro/jobs/download-pro-plugins.py @@ -96,10 +96,11 @@ try: db = Database(LOGGER, sqlalchemy_string=getenv("DATABASE_URI")) db_metadata = db.get_metadata() current_date = datetime.now() + pro_license_key = getenv("PRO_LICENSE_KEY") - # If we already checked in the last hour, skip the check - if db_metadata["last_pro_check"] and (current_date - db_metadata["last_pro_check"]).seconds < 3500: - LOGGER.info("Skipping the check for BunkerWeb Pro license (already checked in the last hour)") + # If we already checked today, skip the check + if pro_license_key == db_metadata["pro_license_key"] and db_metadata["last_pro_check"] and current_date.day == db_metadata["last_pro_check"].day: + LOGGER.info("Skipping the check for BunkerWeb Pro license (already checked today)") sys_exit(0) LOGGER.info("Checking BunkerWeb Pro license key...") @@ -113,21 +114,24 @@ try: headers = {"User-Agent": f"BunkerWeb/{data['version']}"} default_metadata = { "is_pro": False, + "pro_license_key": None, "pro_expire": None, "pro_status": "invalid", "pro_overlapped": False, "pro_services": 0, + "last_pro_check": current_date, } metadata = {} - pro_license_key = getenv("PRO_LICENSE_KEY") error = False temp_dir = TMP_DIR.joinpath(str(uuid4())) temp_dir.mkdir(parents=True, exist_ok=True) if pro_license_key: + default_metadata["pro_license_key"] = (pro_license_key := pro_license_key.strip()) + LOGGER.info("BunkerWeb Pro license provided, checking if it's valid...") - headers["Authorization"] = f"Bearer {pro_license_key.strip()}" + headers["Authorization"] = f"Bearer {pro_license_key}" resp = get(f"{API_ENDPOINT}/pro/status", headers=headers, json=data, timeout=5, allow_redirects=True) if resp.status_code == 403: @@ -153,8 +157,8 @@ try: metadata["pro_overlapped"] = True metadata["is_pro"] = metadata["pro_status"] == "active" - metadata = metadata or default_metadata.copy() - db.set_pro_metadata(metadata | {"last_pro_check": current_date}) + metadata = default_metadata | metadata + db.set_pro_metadata(metadata) if metadata["is_pro"] != db_metadata["is_pro"]: clean_pro_plugins(db) @@ -171,7 +175,7 @@ try: resp_data = resp.json() if resp_data.get("action") == "clean": metadata = default_metadata.copy() - db.set_pro_metadata(metadata | {"last_pro_check": current_date}) + db.set_pro_metadata(metadata) clean_pro_plugins(db) elif resp.headers.get("Content-Type", "") != "application/octet-stream": LOGGER.error(f"Got unexpected content type: {resp.headers.get('Content-Type', 'missing')} from {API_ENDPOINT}/pro") @@ -192,11 +196,6 @@ try: message = "No BunkerWeb Pro license key provided" LOGGER.warning(f"{message}, only checking if there are new or updated preview versions of Pro plugins...") - # If we already checked in the last day, skip the check (only for preview versions) - if not db_metadata["is_pro"] and db_metadata["last_pro_check"] and (current_date - db_metadata["last_pro_check"]).days < 1: - LOGGER.info("Skipping the check for BunkerWeb Pro preview plugins (already checked in the last day)") - sys_exit(0) - resp = get(f"{PREVIEW_ENDPOINT}/v{data['version']}.zip", timeout=5, allow_redirects=True) if resp.status_code == 404: diff --git a/src/common/core/pro/plugin.json b/src/common/core/pro/plugin.json index 2a78ddcb5..93371bd62 100644 --- a/src/common/core/pro/plugin.json +++ b/src/common/core/pro/plugin.json @@ -19,7 +19,7 @@ { "name": "download-pro-plugins", "file": "download-pro-plugins.py", - "every": "hour", + "every": "day", "reload": true } ] diff --git a/src/common/db/Database.py b/src/common/db/Database.py index 4a5594c84..a94ad5e57 100644 --- a/src/common/db/Database.py +++ b/src/common/db/Database.py @@ -236,7 +236,7 @@ class Database: return "" - def set_pro_metadata(self, data: Dict[Literal["is_pro", "pro_expire", "pro_status", "pro_overlapped", "pro_services"], Any] = {}) -> str: + def set_pro_metadata(self, data: Dict[Literal["is_pro", "pro_license_key", "pro_expire", "pro_status", "pro_overlapped", "pro_services"], Any] = {}) -> str: """Set the pro metadata values""" with self.__db_session() as session: try: @@ -312,6 +312,7 @@ class Database: "integration": "unknown", "database_version": "Unknown", "is_pro": "no", + "pro_license_key": "", "pro_expire": None, "pro_services": 0, "pro_overlapped": False, @@ -331,6 +332,7 @@ class Database: Metadata.version, Metadata.integration, Metadata.is_pro, + Metadata.pro_license_key, Metadata.pro_expire, Metadata.pro_services, Metadata.pro_overlapped, @@ -346,6 +348,7 @@ class Database: "version": metadata.version, "integration": metadata.integration, "is_pro": metadata.is_pro, + "pro_license_key": metadata.pro_license_key or "", "pro_expire": metadata.pro_expire, "pro_services": metadata.pro_services, "pro_overlapped": metadata.pro_overlapped, diff --git a/src/common/db/model.py b/src/common/db/model.py index 67c649d58..c59ae120e 100644 --- a/src/common/db/model.py +++ b/src/common/db/model.py @@ -216,6 +216,7 @@ class Metadata(Base): id = Column(Integer, primary_key=True, default=1) is_initialized = Column(Boolean, nullable=False) is_pro = Column(Boolean, default=False, nullable=False) + pro_license_key = Column(String(256), nullable=True) pro_expire = Column(DateTime, nullable=True) pro_status = Column(PRO_STATUS_ENUM, default="invalid", nullable=False) pro_services = Column(Integer, default=0, nullable=False)