LibreTranslate/libretranslate/api_keys.py

113 lines
3.7 KiB
Python
Raw Permalink Normal View History

import os
import sqlite3
import uuid
2023-07-09 10:29:11 +00:00
2022-06-21 18:57:32 +00:00
import requests
from expiringdict import ExpiringDict
2023-07-09 10:29:11 +00:00
from libretranslate.default_values import DEFAULT_ARGUMENTS as DEFARGS
DEFAULT_DB_PATH = DEFARGS['API_KEYS_DB_PATH']
2021-05-18 03:41:02 +00:00
class Database:
2021-05-18 03:41:02 +00:00
def __init__(self, db_path=DEFAULT_DB_PATH, max_cache_len=1000, max_cache_age=30):
# Legacy check - this can be removed at some point in the near future
if os.path.isfile("api_keys.db") and not os.path.isfile("db/api_keys.db"):
2023-07-09 10:29:11 +00:00
print("Migrating {} to {}".format("api_keys.db", "db/api_keys.db"))
try:
os.rename("api_keys.db", "db/api_keys.db")
except Exception as e:
print(str(e))
db_dir = os.path.dirname(db_path)
2023-07-09 10:29:11 +00:00
if db_dir != '' and not os.path.exists(db_dir):
os.makedirs(db_dir)
self.db_path = db_path
self.cache = ExpiringDict(max_len=max_cache_len, max_age_seconds=max_cache_age)
# Make sure to do data synchronization on writes!
self.c = sqlite3.connect(db_path, check_same_thread=False)
2021-05-18 03:41:02 +00:00
self.c.execute(
"""CREATE TABLE IF NOT EXISTS api_keys (
"api_key" TEXT NOT NULL,
"req_limit" INTEGER NOT NULL,
2024-01-21 17:48:42 +00:00
"char_limit" INTEGER DEFAULT NULL,
PRIMARY KEY("api_key")
2021-05-18 03:41:02 +00:00
);"""
)
2024-01-21 17:48:42 +00:00
# Schema/upgrade checks
schema = self.c.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='api_keys';").fetchone()[0]
if '"char_limit" INTEGER DEFAULT NULL' not in schema:
self.c.execute('ALTER TABLE api_keys ADD COLUMN "char_limit" INTEGER DEFAULT NULL;')
def lookup(self, api_key):
2024-01-21 17:48:42 +00:00
val = self.cache.get(api_key)
if val is None:
# DB Lookup
2021-05-18 03:41:02 +00:00
stmt = self.c.execute(
2024-01-21 17:48:42 +00:00
"SELECT req_limit, char_limit FROM api_keys WHERE api_key = ?", (api_key,)
2021-05-18 03:41:02 +00:00
)
row = stmt.fetchone()
if row is not None:
2024-01-21 17:48:42 +00:00
self.cache[api_key] = row
val = row
else:
self.cache[api_key] = False
2024-01-21 17:48:42 +00:00
val = False
2021-05-18 03:41:02 +00:00
2024-01-21 17:48:42 +00:00
if isinstance(val, bool):
val = None
2021-05-18 03:41:02 +00:00
2024-01-21 17:48:42 +00:00
return val
2021-05-18 03:41:02 +00:00
2024-01-21 17:48:42 +00:00
def add(self, req_limit, api_key="auto", char_limit=None):
if api_key == "auto":
api_key = str(uuid.uuid4())
2024-01-21 17:48:42 +00:00
if char_limit == 0:
char_limit = None
self.remove(api_key)
2021-05-18 03:41:02 +00:00
self.c.execute(
2024-01-21 17:48:42 +00:00
"INSERT INTO api_keys (api_key, req_limit, char_limit) VALUES (?, ?, ?)",
(api_key, req_limit, char_limit),
2021-05-18 03:41:02 +00:00
)
self.c.commit()
2024-01-21 17:48:42 +00:00
return (api_key, req_limit, char_limit)
2021-05-18 03:41:02 +00:00
def remove(self, api_key):
2021-05-18 03:41:02 +00:00
self.c.execute("DELETE FROM api_keys WHERE api_key = ?", (api_key,))
self.c.commit()
return api_key
def all(self):
2024-01-21 17:48:42 +00:00
row = self.c.execute("SELECT api_key, req_limit, char_limit FROM api_keys")
2021-05-18 03:41:02 +00:00
return row.fetchall()
2022-06-21 18:57:32 +00:00
class RemoteDatabase:
def __init__(self, url, max_cache_len=1000, max_cache_age=600):
self.url = url
self.cache = ExpiringDict(max_len=max_cache_len, max_age_seconds=max_cache_age)
def lookup(self, api_key):
2024-01-21 17:48:42 +00:00
val = self.cache.get(api_key)
if val is None:
2022-06-21 18:57:32 +00:00
try:
r = requests.post(self.url, data={'api_key': api_key}, timeout=60)
2022-06-21 18:57:32 +00:00
res = r.json()
except Exception as e:
print("Cannot authenticate API key: " + str(e))
return None
2022-06-21 18:57:32 +00:00
2024-01-21 18:17:07 +00:00
if res.get('error') is not None:
return None
req_limit = res.get('req_limit', None)
char_limit = res.get('char_limit', None)
2024-01-21 17:48:42 +00:00
self.cache[api_key] = (req_limit, char_limit)
2022-06-21 18:57:32 +00:00
2024-01-21 17:48:42 +00:00
return val