import os import time import asyncio import json from dotenv import load_dotenv from motor.motor_asyncio import AsyncIOMotorClient load_dotenv() MONGO_URI = os.getenv("MONGO_URI") async def log_rated_key(key: str) -> None: """Logs a key that has been rate limited to the database.""" client = AsyncIOMotorClient(MONGO_URI) scheme = { "key": key, "timestamp_added": int(time.time()) } collection = client['Liabilities']['rate-limited-keys'] await collection.insert_one(scheme) async def key_is_rated(key: str) -> bool: """Checks if a key is rate limited.""" client = AsyncIOMotorClient(MONGO_URI) collection = client['Liabilities']['rate-limited-keys'] query = { "key": key } result = await collection.find_one(query) return result is not None async def cached_key_is_rated(key: str) -> bool: path = os.path.join(os.getcwd(), "cache", "rate_limited_keys.json") with open(path, "r") as file: keys = json.load(file) return key in keys async def remove_rated_keys() -> None: """Removes all keys that have been rate limited for more than a day.""" a_day = 86400 client = AsyncIOMotorClient(MONGO_URI) collection = client['Liabilities']['rate-limited-keys'] keys = await collection.find().to_list(length=None) marked_for_removal = [] for key in keys: if int(time.time()) - key['timestamp_added'] > a_day: marked_for_removal.append(key['_id']) query = { "_id": { "$in": marked_for_removal } } await collection.delete_many(query) async def cache_all_keys() -> None: """Clones all keys from the database to the cache.""" client = AsyncIOMotorClient(MONGO_URI) collection = client['Liabilities']['rate-limited-keys'] keys = await collection.find().to_list(length=None) keys = [key['key'] for key in keys] path = os.path.join(os.getcwd(), "cache", "rate_limited_keys.json") with open(path, "w") as file: json.dump(keys, file) if __name__ == "__main__": asyncio.run(remove_rated_keys())