Add shared storage option

This commit is contained in:
Piero Toffanin 2023-03-09 22:00:27 -05:00
parent e00f3af7db
commit ab27dcbaf3
5 changed files with 204 additions and 10 deletions

View file

@ -270,14 +270,12 @@ def create_app(args):
):
need_key = True
# TODO: find a way to send a "refresh" error key?
if need_key:
description = _("Please contact the server operator to get an API key")
if args.get_api_key_link:
description = _("Visit %(url)s to get an API key", url=args.get_api_key_link)
abort(
403,
400,
description=description,
)
return f(*a, **kw)

42
libretranslate/secret.py Normal file
View file

@ -0,0 +1,42 @@
import atexit
import random
import string
from multiprocessing import Value
from libretranslate.storage import get_storage
from apscheduler.schedulers.background import BackgroundScheduler
setup_secrets = Value('b', False)
def generate_secret():
return ''.join(random.choices(string.ascii_uppercase + string.digits, k=7))
def rotate_secrets():
s = get_storage()
secret_1 = s.get_str("secret_1")
s.set_str("secret_0", secret_1)
s.set_str("secret_1", generate_secret())
def secret_match(secret):
s = get_storage()
return secret == s.get_str("secret_0") or secret == s.get_str("secret_1")
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
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())

147
libretranslate/storage.py Normal file
View file

@ -0,0 +1,147 @@
import redis
storage = None
def get_storage():
return storage
class Storage:
def set_bool(self, key, value):
raise Exception("not implemented")
def get_bool(self, key):
raise Exception("not implemented")
def set_int(self, key, value):
raise Exception("not implemented")
def get_int(self, key):
raise Exception("not implemented")
def set_str(self, key, value):
raise Exception("not implemented")
def get_str(self, key):
raise Exception("not implemented")
def set_hash_value(self, ns, key, value):
raise Exception("not implemented")
def get_hash_int(self, ns, key):
raise Exception("not implemented")
def inc_hash_int(self, ns, key):
raise Exception("not implemented")
def dec_hash_int(self, ns, key):
raise Exception("not implemented")
def get_hash_keys(self, ns):
raise Exception("not implemented")
def del_hash(self, ns, key):
raise Exception("not implemented")
class MemoryStorage(Storage):
def __init__(self):
self.store = {}
def set_bool(self, key, value):
self.store[key] = bool(value)
def get_bool(self, key):
return bool(self.store[key])
def set_int(self, key, value):
self.store[key] = int(value)
def get_int(self, key):
return int(self.store.get(key, 0))
def set_str(self, key, value):
self.store[key] = value
def get_str(self, key):
return str(self.store.get(key, ""))
def get_hash_int(self, ns, key):
d = self.store.get(ns, {})
return int(d.get(key, 0))
def inc_hash_int(self, ns, key):
if ns not in self.store:
self.store[ns] = {}
if key not in self.store[ns]:
self.store[ns][key] = 0
else:
self.store[ns][key] += 1
def dec_hash_int(self, ns, key):
if ns not in self.store:
self.store[ns] = {}
if key not in self.store[ns]:
self.store[ns][key] = 0
else:
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()]
def del_hash(self, ns, key):
del self.store[ns][key]
class RedisStorage(Storage):
def __init__(self, redis_uri):
self.conn = redis.from_url(redis_uri)
self.conn.ping()
def set_bool(self, key, value):
self.conn.set(key, "1" if value else "0")
def get_bool(self, key):
return bool(self.conn.get(key))
def set_int(self, key, value):
self.conn.set(key, str(value))
def get_int(self, key):
v = self.conn.get(key)
if v is None:
return 0
else:
return v
def set_str(self, key, value):
self.conn.set(key, value)
def get_str(self, key):
v = self.conn.get(key)
if v is None:
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 inc_hash_int(self, ns, key):
return int(self.conn.hincrby(ns, key))
def dec_hash_int(self, ns, key):
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()]
def del_hash(self, ns, key):
conn.hdel(ns, key)
def setup(storage_uri):
global storage
if storage_uri.startswith("memory://"):
storage = MemoryStorage()
elif storage_uri.startswith("redis://"):
storage = RedisStorage(storage_uri)
else:
raise Exception("Invalid storage URI: " + storage_uri)
return storage

View file

@ -243,9 +243,8 @@ document.addEventListener('DOMContentLoaded', function(){
request.onload = function() {
try{
{% if api_secret != "" %}
if (this.status === 403){
//window.location.reload(true);
//return;
if (this.status === 400){
if (self.refreshOnce()) return;
}
{% endif %}
@ -362,6 +361,15 @@ document.addEventListener('DOMContentLoaded', function(){
this.translatedFileUrl = false;
this.loadingFileTranslation = false;
},
refreshOnce: function(){
var lastRefreshed = parseInt(localStorage.getItem("refreshed") || 0);
var now = new Date().getTime();
if (now - lastRefreshed > 1000 * 60 * 1){
localStorage.setItem("refreshed", now);
window.location.reload();
return true;
}
},
translateFile: function(e) {
e.preventDefault();
@ -383,9 +391,8 @@ document.addEventListener('DOMContentLoaded', function(){
if (translateFileRequest.readyState === 4 && translateFileRequest.status === 200) {
try{
{% if api_secret != "" %}
if (this.status === 403){
window.location.reload(true);
return;
if (this.status === 400){
if (self.refreshOnce()) return;
}
{% endif %}
self.loadingFileTranslation = false;