From 8f8087f8ae3b073ce0fc19c01315472000176b2b Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Tue, 21 Jun 2022 14:57:32 -0400
Subject: [PATCH] Add --api-keys-remote

---
 README.md             |  1 +
 app/api_keys.py       | 27 ++++++++++++++++++++++++++-
 app/app.py            |  9 +++++++--
 app/default_values.py |  5 +++++
 app/main.py           |  6 ++++++
 requirements.txt      |  2 ++
 6 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 706cb24..bdb4eb3 100644
--- a/README.md
+++ b/README.md
@@ -190,6 +190,7 @@ docker-compose -f docker-compose.cuda.yml up -d --build
 | --frontend-language-target  | Set frontend default language - target                                                                      | `es`          | LT_FRONTEND_LANGUAGE_TARGET  |
 | --frontend-timeout          | Set frontend translation timeout                                                                            | `500`         | LT_FRONTEND_TIMEOUT          |
 | --api-keys                  | Enable API keys database for per-user rate limits lookup                                                    | `Don't use API keys` | LT_API_KEYS                  |
+| --api-keys-remote           | Use this remote endpoint to query for valid API keys instead of using the local database                    | `Use local API key database` | LT_API_KEYS_REMOTE                  |
 | --require-api-key-origin    | Require use of an API key for programmatic access to the API, unless the request origin matches this domain | `No restrictions on domain origin` | LT_REQUIRE_API_KEY_ORIGIN    |
 | --load-only                 | Set available languages                                                                                     | `all from argostranslate`    | LT_LOAD_ONLY                 |
 | --suggestions               | Allow user suggestions                                                                                      | `false`    | LT_SUGGESTIONS               |
diff --git a/app/api_keys.py b/app/api_keys.py
index 905047b..eb2d496 100644
--- a/app/api_keys.py
+++ b/app/api_keys.py
@@ -1,6 +1,6 @@
 import sqlite3
 import uuid
-
+import requests
 from expiringdict import ExpiringDict
 
 DEFAULT_DB_PATH = "api_keys.db"
@@ -61,3 +61,28 @@ class Database:
     def all(self):
         row = self.c.execute("SELECT api_key, req_limit FROM api_keys")
         return row.fetchall()
+
+
+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):
+        req_limit = self.cache.get(api_key)
+        if req_limit is None:
+            try:
+                r = requests.post(self.url, data={'api_key': api_key})
+                res = r.json()
+            except Exception as e:
+                print("Cannot authenticate API key: " + str(e))
+                return False
+
+            if res.get('error', None) is None:
+                req_limit = res.get('req_limit', None)
+                self.cache[api_key] = req_limit
+            else:
+                req_limit = False
+                self.cache[api_key] = False
+
+        return req_limit
diff --git a/app/app.py b/app/app.py
index 935d7a0..cba22a2 100644
--- a/app/app.py
+++ b/app/app.py
@@ -17,7 +17,7 @@ from werkzeug.utils import secure_filename
 from app import flood, remove_translated_files, security
 from app.language import detect_languages, transliterate
 
-from .api_keys import Database
+from .api_keys import Database, RemoteDatabase
 from .suggestions import Database as SuggestionsDatabase
 
 
@@ -146,7 +146,12 @@ def create_app(args):
     api_keys_db = None
 
     if args.req_limit > 0 or args.api_keys or args.daily_req_limit > 0:
-        api_keys_db = Database() if args.api_keys else None
+        api_keys_db = None
+        if args.api_keys:
+            if args.api_keys_remote:
+                api_keys_db = RemoteDatabase(args.api_keys_remote)
+            else:
+                api_keys_db = Database()
 
         from flask_limiter import Limiter
 
diff --git a/app/default_values.py b/app/default_values.py
index 9161aa6..f8043b5 100644
--- a/app/default_values.py
+++ b/app/default_values.py
@@ -105,6 +105,11 @@ _default_options_objects = [
         'name': 'API_KEYS',
         'default_value': False,
         'value_type': 'bool'
+    },    
+    {
+        'name': 'API_KEYS_REMOTE',
+        'default_value': '',
+        'value_type': 'str'
     },
     {
         'name': 'REQUIRE_API_KEY_ORIGIN',
diff --git a/app/main.py b/app/main.py
index 1545a77..9e44050 100644
--- a/app/main.py
+++ b/app/main.py
@@ -89,6 +89,12 @@ def get_args():
         action="store_true",
         help="Enable API keys database for per-user rate limits lookup",
     )
+    parser.add_argument(
+        "--api-keys-remote",
+        default=DEFARGS['API_KEYS_REMOTE'],
+        type=str,
+        help="Use this remote endpoint to query for valid API keys instead of using the local database",
+    )
     parser.add_argument(
         "--require-api-key-origin",
         type=str,
diff --git a/requirements.txt b/requirements.txt
index 4f605d5..9ea3949 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,3 +15,5 @@ translatehtml==1.5.2
 argos-translate-files==1.1.0
 itsdangerous==2.1.2
 Werkzeug==2.1.2
+requests==2.28.0
+