46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
from my_modules.app.constens import SECRET_KEY
|
|
|
|
import hmac, hashlib, base64, secrets, string, time
|
|
from datetime import datetime, timezone
|
|
|
|
def base64url_encode(data: bytes) -> str:
|
|
return base64.urlsafe_b64encode(data).decode().rstrip("=")
|
|
|
|
def base64url_decode(data: str) -> bytes:
|
|
padding = '=' * (-len(data) % 4)
|
|
return base64.urlsafe_b64decode(data + padding)
|
|
|
|
def generate_short_id(length=11):
|
|
alphabet = string.ascii_letters + string.digits + '_-'
|
|
while True:
|
|
token = ''.join(secrets.choice(alphabet) for _ in range(length))
|
|
if not token.startswith('-'):
|
|
return token
|
|
|
|
def generate_signed_url(file_id: str) -> str:
|
|
# signature based only on the file_id
|
|
sig = hmac.new(SECRET_KEY, file_id.encode(), hashlib.sha256).digest()
|
|
token = base64url_encode(sig)
|
|
return f"-{file_id}?sig={token}"
|
|
|
|
def verify_signed_url(file_id: str, token: str, file_expiration: int) -> bool:
|
|
# check both the signature and the file's stored expiration time
|
|
expected_sig = hmac.new(SECRET_KEY, file_id.encode(), hashlib.sha256).digest()
|
|
valid_sig = hmac.compare_digest(base64url_encode(expected_sig), token)
|
|
not_expired = file_expiration >= time.time()
|
|
return valid_sig and not_expired
|
|
|
|
def is_expired(expires_at):
|
|
if not expires_at:
|
|
return False
|
|
if expires_at.tzinfo is None:
|
|
expires_at = expires_at.replace(tzinfo=timezone.utc)
|
|
else:
|
|
expires_at = expires_at.astimezone(timezone.utc)
|
|
return expires_at <= datetime.now(timezone.utc)
|
|
|
|
if __name__ == "__main__":
|
|
file_id = generate_short_id()
|
|
url = generate_signed_url(file_id)
|
|
print(url)
|