From 6a231d974bd17e800b0bff93d9e2b858ad3678ea Mon Sep 17 00:00:00 2001 From: Piero Toffanin Date: Thu, 9 Mar 2023 23:07:12 -0500 Subject: [PATCH] Testing, fixes --- libretranslate/app.py | 13 ++++++------ libretranslate/flood.py | 28 +++++--------------------- libretranslate/scheduler.py | 23 +++++++++++++++++++++ libretranslate/secret.py | 20 ++++++------------- libretranslate/storage.py | 21 ++++++++++++++----- scripts/gunicorn_conf.py | 40 ++++++++++++++++++++++++++++++++++++- 6 files changed, 96 insertions(+), 49 deletions(-) create mode 100644 libretranslate/scheduler.py diff --git a/libretranslate/app.py b/libretranslate/app.py index 47ad733..ea4c0b6 100644 --- a/libretranslate/app.py +++ b/libretranslate/app.py @@ -21,7 +21,7 @@ from werkzeug.exceptions import HTTPException from werkzeug.http import http_date from flask_babel import Babel -from libretranslate import flood, secret, remove_translated_files, security, storage +from libretranslate import scheduler, flood, secret, remove_translated_files, security, storage from libretranslate.language import detect_languages, improve_translation_formatting from libretranslate.locales import (_, _lazy, get_available_locales, get_available_locale_codes, gettext_escaped, gettext_html, lazy_swag, get_alternate_locale_links) @@ -204,11 +204,12 @@ def create_app(args): limiter = Limiter() - if args.req_flood_threshold > 0: - flood.setup(args.req_flood_threshold) - if args.api_keys and args.require_api_key_secret: - secret.setup() - + if not "gunicorn" in os.environ.get("SERVER_SOFTWARE", ""): + # Gunicorn starts the scheduler in the master process + scheduler.setup(args) + + flood.setup(args) + secret.setup(args) measure_request = None gauge_request = None diff --git a/libretranslate/flood.py b/libretranslate/flood.py index 6990429..63cd392 100644 --- a/libretranslate/flood.py +++ b/libretranslate/flood.py @@ -1,11 +1,5 @@ -import atexit -from multiprocessing import Value - from libretranslate.storage import get_storage -from apscheduler.schedulers.background import BackgroundScheduler - -setup_scheduler = Value('b', False) active = False threshold = -1 @@ -20,30 +14,18 @@ def forgive_banned(): if banned[ip] <= 0: clear_list.append(ip) else: - banned[ip] = min(threshold, banned[ip]) - 1 + s.set_hash_int("banned", ip, min(threshold, banned[ip]) - 1) for ip in clear_list: s.del_hash("banned", ip) -def setup(violations_threshold=100): +def setup(args): global active global threshold - active = True - threshold = violations_threshold - - # Only setup the scheduler and secrets on one process - if not setup_scheduler.value: - setup_scheduler.value = True - - scheduler = BackgroundScheduler() - scheduler.add_job(func=forgive_banned, trigger="interval", minutes=30) - - scheduler.start() - - # Shut down the scheduler when exiting the app - atexit.register(lambda: scheduler.shutdown()) - + if args.req_flood_threshold > 0: + active = True + threshold = args.req_flood_threshold def report(request_ip): if active: diff --git a/libretranslate/scheduler.py b/libretranslate/scheduler.py new file mode 100644 index 0000000..3300095 --- /dev/null +++ b/libretranslate/scheduler.py @@ -0,0 +1,23 @@ +import atexit +from apscheduler.schedulers.background import BackgroundScheduler +scheduler = None + +def setup(args): + from libretranslate.flood import forgive_banned + from libretranslate.secret import rotate_secrets + + global scheduler + + if scheduler is None: + scheduler = BackgroundScheduler() + + if args.req_flood_threshold > 0: + scheduler.add_job(func=forgive_banned, trigger="interval", minutes=10) + + if args.api_keys and args.require_api_key_secret: + scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30) + + scheduler.start() + + # Shut down the scheduler when exiting the app + atexit.register(lambda: scheduler.shutdown()) \ No newline at end of file diff --git a/libretranslate/secret.py b/libretranslate/secret.py index 237dc84..1651e8b 100644 --- a/libretranslate/secret.py +++ b/libretranslate/secret.py @@ -1,7 +1,7 @@ import atexit import random import string -from multiprocessing import Value +from multiprocessing.dummy import Value from libretranslate.storage import get_storage from apscheduler.schedulers.background import BackgroundScheduler @@ -16,6 +16,9 @@ def rotate_secrets(): secret_1 = s.get_str("secret_1") s.set_str("secret_0", secret_1) s.set_str("secret_1", generate_secret()) + print(s.get_str("secret_0")) + print(s.get_str("secret_1")) + def secret_match(secret): s = get_storage() @@ -24,19 +27,8 @@ def secret_match(secret): def get_current_secret(): return get_storage().get_str("secret_1") -def setup(): - # Only setup the scheduler and secrets on one process - if not setup_secrets.value: - setup_secrets.value = True - +def setup(args): + if args.api_keys and args.require_api_key_secret: s = get_storage() s.set_str("secret_0", generate_secret()) s.set_str("secret_1", generate_secret()) - - scheduler = BackgroundScheduler() - scheduler.add_job(func=rotate_secrets, trigger="interval", minutes=30) - - scheduler.start() - - # Shut down the scheduler when exiting the app - atexit.register(lambda: scheduler.shutdown()) diff --git a/libretranslate/storage.py b/libretranslate/storage.py index 409c79c..7ffe849 100644 --- a/libretranslate/storage.py +++ b/libretranslate/storage.py @@ -20,7 +20,7 @@ class Storage: def get_str(self, key): raise Exception("not implemented") - def set_hash_value(self, ns, key, value): + def set_hash_int(self, ns, key, value): raise Exception("not implemented") def get_hash_int(self, ns, key): raise Exception("not implemented") @@ -56,6 +56,11 @@ class MemoryStorage(Storage): def get_str(self, key): return str(self.store.get(key, "")) + def set_hash_int(self, ns, key, value): + if ns not in self.store: + self.store[ns] = {} + self.store[ns][key] = int(value) + def get_hash_int(self, ns, key): d = self.store.get(ns, {}) return int(d.get(key, 0)) @@ -79,7 +84,10 @@ class MemoryStorage(Storage): self.store[ns][key] -= 1 def get_all_hash_int(self, ns): - return [{str(k): int(v)} for k,v in self.store[ns].items()] + if ns in self.store: + return [{str(k): int(v)} for k,v in self.store[ns].items()] + else: + return [] def del_hash(self, ns, key): del self.store[ns][key] @@ -115,13 +123,16 @@ class RedisStorage(Storage): return "" else: return v.decode('utf-8') - + def get_hash_int(self, ns, key): v = self.conn.hget(ns, key) if v is None: return 0 else: return int(v) + + def set_hash_int(self, ns, key, value): + self.conn.hset(ns, key, value) def inc_hash_int(self, ns, key): return int(self.conn.hincrby(ns, key)) @@ -130,10 +141,10 @@ class RedisStorage(Storage): return int(self.conn.hincrby(ns, key, -1)) def get_all_hash_int(self, ns): - return [{k.decode("utf-8"): int(v)} for k,v in self.conn.hgetall(ns).items()] + return {k.decode("utf-8"): int(v) for k,v in self.conn.hgetall(ns).items()} def del_hash(self, ns, key): - conn.hdel(ns, key) + self.conn.hdel(ns, key) def setup(storage_uri): global storage diff --git a/scripts/gunicorn_conf.py b/scripts/gunicorn_conf.py index 6ec7e40..845da5b 100644 --- a/scripts/gunicorn_conf.py +++ b/scripts/gunicorn_conf.py @@ -1,4 +1,42 @@ from prometheus_client import multiprocess +import re +import sys def child_exit(server, worker): - multiprocess.mark_process_dead(worker.pid) \ No newline at end of file + multiprocess.mark_process_dead(worker.pid) + +def on_starting(server): + # Parse command line arguments + proc_name = server.cfg.default_proc_name + kwargs = {} + if proc_name.startswith("wsgi:app"): + str_args = re.sub('wsgi:app\s*\(\s*(.*)\s*\)', '\\1', proc_name).strip().split(",") + for a in str_args: + if "=" in a: + k,v = a.split("=") + k = k.strip() + v = v.strip() + + if v.lower() in ["true", "false"]: + v = v.lower() == "true" + elif v[0] == '"': + v = v[1:-1] + kwargs[k] = v + + from libretranslate.main import get_args + sys.argv = ['--wsgi'] + for k in kwargs: + ck = k.replace("_", "-") + if isinstance(kwargs[k], bool) and kwargs[k]: + sys.argv.append("--" + ck) + else: + sys.argv.append("--" + ck) + sys.argv.append(kwargs[k]) + + args = get_args() + + from libretranslate import storage, scheduler, flood, secret + storage.setup(args.shared_storage) + scheduler.setup(args) + flood.setup(args) + secret.setup(args) \ No newline at end of file