|
@@ -1,18 +1,23 @@
|
|
import json
|
|
import json
|
|
|
|
+import logging
|
|
import random
|
|
import random
|
|
import re
|
|
import re
|
|
import string
|
|
import string
|
|
import subprocess
|
|
import subprocess
|
|
|
|
+import time
|
|
import uuid
|
|
import uuid
|
|
from collections.abc import Generator
|
|
from collections.abc import Generator
|
|
from datetime import datetime
|
|
from datetime import datetime
|
|
from hashlib import sha256
|
|
from hashlib import sha256
|
|
-from typing import Union
|
|
|
|
|
|
+from typing import Any, Optional, Union
|
|
from zoneinfo import available_timezones
|
|
from zoneinfo import available_timezones
|
|
|
|
|
|
-from flask import Response, stream_with_context
|
|
|
|
|
|
+from flask import Response, current_app, stream_with_context
|
|
from flask_restful import fields
|
|
from flask_restful import fields
|
|
|
|
|
|
|
|
+from extensions.ext_redis import redis_client
|
|
|
|
+from models.account import Account
|
|
|
|
+
|
|
|
|
|
|
def run(script):
|
|
def run(script):
|
|
return subprocess.getstatusoutput('source /root/.bashrc && ' + script)
|
|
return subprocess.getstatusoutput('source /root/.bashrc && ' + script)
|
|
@@ -46,12 +51,12 @@ def uuid_value(value):
|
|
error = ('{value} is not a valid uuid.'
|
|
error = ('{value} is not a valid uuid.'
|
|
.format(value=value))
|
|
.format(value=value))
|
|
raise ValueError(error)
|
|
raise ValueError(error)
|
|
-
|
|
|
|
|
|
+
|
|
def alphanumeric(value: str):
|
|
def alphanumeric(value: str):
|
|
# check if the value is alphanumeric and underlined
|
|
# check if the value is alphanumeric and underlined
|
|
if re.match(r'^[a-zA-Z0-9_]+$', value):
|
|
if re.match(r'^[a-zA-Z0-9_]+$', value):
|
|
return value
|
|
return value
|
|
-
|
|
|
|
|
|
+
|
|
raise ValueError(f'{value} is not a valid alphanumeric value')
|
|
raise ValueError(f'{value} is not a valid alphanumeric value')
|
|
|
|
|
|
def timestamp_value(timestamp):
|
|
def timestamp_value(timestamp):
|
|
@@ -163,3 +168,97 @@ def compact_generate_response(response: Union[dict, Generator]) -> Response:
|
|
|
|
|
|
return Response(stream_with_context(generate()), status=200,
|
|
return Response(stream_with_context(generate()), status=200,
|
|
mimetype='text/event-stream')
|
|
mimetype='text/event-stream')
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+class TokenManager:
|
|
|
|
+
|
|
|
|
+ @classmethod
|
|
|
|
+ def generate_token(cls, account: Account, token_type: str, additional_data: dict = None) -> str:
|
|
|
|
+ old_token = cls._get_current_token_for_account(account.id, token_type)
|
|
|
|
+ if old_token:
|
|
|
|
+ if isinstance(old_token, bytes):
|
|
|
|
+ old_token = old_token.decode('utf-8')
|
|
|
|
+ cls.revoke_token(old_token, token_type)
|
|
|
|
+
|
|
|
|
+ token = str(uuid.uuid4())
|
|
|
|
+ token_data = {
|
|
|
|
+ 'account_id': account.id,
|
|
|
|
+ 'email': account.email,
|
|
|
|
+ 'token_type': token_type
|
|
|
|
+ }
|
|
|
|
+ if additional_data:
|
|
|
|
+ token_data.update(additional_data)
|
|
|
|
+
|
|
|
|
+ expiry_hours = current_app.config[f'{token_type.upper()}_TOKEN_EXPIRY_HOURS']
|
|
|
|
+ token_key = cls._get_token_key(token, token_type)
|
|
|
|
+ redis_client.setex(
|
|
|
|
+ token_key,
|
|
|
|
+ expiry_hours * 60 * 60,
|
|
|
|
+ json.dumps(token_data)
|
|
|
|
+ )
|
|
|
|
+
|
|
|
|
+ cls._set_current_token_for_account(account.id, token, token_type, expiry_hours)
|
|
|
|
+ return token
|
|
|
|
+
|
|
|
|
+ @classmethod
|
|
|
|
+ def _get_token_key(cls, token: str, token_type: str) -> str:
|
|
|
|
+ return f'{token_type}:token:{token}'
|
|
|
|
+
|
|
|
|
+ @classmethod
|
|
|
|
+ def revoke_token(cls, token: str, token_type: str):
|
|
|
|
+ token_key = cls._get_token_key(token, token_type)
|
|
|
|
+ redis_client.delete(token_key)
|
|
|
|
+
|
|
|
|
+ @classmethod
|
|
|
|
+ def get_token_data(cls, token: str, token_type: str) -> Optional[dict[str, Any]]:
|
|
|
|
+ key = cls._get_token_key(token, token_type)
|
|
|
|
+ token_data_json = redis_client.get(key)
|
|
|
|
+ if token_data_json is None:
|
|
|
|
+ logging.warning(f"{token_type} token {token} not found with key {key}")
|
|
|
|
+ return None
|
|
|
|
+ token_data = json.loads(token_data_json)
|
|
|
|
+ return token_data
|
|
|
|
+
|
|
|
|
+ @classmethod
|
|
|
|
+ def _get_current_token_for_account(cls, account_id: str, token_type: str) -> Optional[str]:
|
|
|
|
+ key = cls._get_account_token_key(account_id, token_type)
|
|
|
|
+ current_token = redis_client.get(key)
|
|
|
|
+ return current_token
|
|
|
|
+
|
|
|
|
+ @classmethod
|
|
|
|
+ def _set_current_token_for_account(cls, account_id: str, token: str, token_type: str, expiry_hours: int):
|
|
|
|
+ key = cls._get_account_token_key(account_id, token_type)
|
|
|
|
+ redis_client.setex(key, expiry_hours * 60 * 60, token)
|
|
|
|
+
|
|
|
|
+ @classmethod
|
|
|
|
+ def _get_account_token_key(cls, account_id: str, token_type: str) -> str:
|
|
|
|
+ return f'{token_type}:account:{account_id}'
|
|
|
|
+
|
|
|
|
+
|
|
|
|
+class RateLimiter:
|
|
|
|
+ def __init__(self, prefix: str, max_attempts: int, time_window: int):
|
|
|
|
+ self.prefix = prefix
|
|
|
|
+ self.max_attempts = max_attempts
|
|
|
|
+ self.time_window = time_window
|
|
|
|
+
|
|
|
|
+ def _get_key(self, email: str) -> str:
|
|
|
|
+ return f"{self.prefix}:{email}"
|
|
|
|
+
|
|
|
|
+ def is_rate_limited(self, email: str) -> bool:
|
|
|
|
+ key = self._get_key(email)
|
|
|
|
+ current_time = int(time.time())
|
|
|
|
+ window_start_time = current_time - self.time_window
|
|
|
|
+
|
|
|
|
+ redis_client.zremrangebyscore(key, '-inf', window_start_time)
|
|
|
|
+ attempts = redis_client.zcard(key)
|
|
|
|
+
|
|
|
|
+ if attempts and int(attempts) >= self.max_attempts:
|
|
|
|
+ return True
|
|
|
|
+ return False
|
|
|
|
+
|
|
|
|
+ def increment_rate_limit(self, email: str):
|
|
|
|
+ key = self._get_key(email)
|
|
|
|
+ current_time = int(time.time())
|
|
|
|
+
|
|
|
|
+ redis_client.zadd(key, {current_time: current_time})
|
|
|
|
+ redis_client.expire(key, self.time_window * 2)
|