2021-08-23 06:14:40 +00:00
|
|
|
from fastapi import Security, HTTPException, APIRouter, Depends
|
|
|
|
from fastapi.security.api_key import APIKeyQuery, APIKeyCookie, APIKeyHeader, APIKey
|
|
|
|
|
|
|
|
from starlette.status import HTTP_403_FORBIDDEN
|
|
|
|
|
|
|
|
import logging
|
|
|
|
import uuid
|
2022-03-13 04:50:33 +00:00
|
|
|
|
|
|
|
from core.storage import Storage
|
2021-08-23 06:14:40 +00:00
|
|
|
|
|
|
|
router = APIRouter()
|
|
|
|
logger = logging.getLogger("security")
|
2022-03-13 04:50:33 +00:00
|
|
|
storage = Storage("security")
|
2021-08-23 06:14:40 +00:00
|
|
|
|
|
|
|
API_KEY_NAME = "access_token"
|
|
|
|
|
|
|
|
api_key_query = APIKeyQuery(name=API_KEY_NAME, auto_error=False)
|
|
|
|
api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
|
|
|
|
api_key_cookie = APIKeyCookie(name=API_KEY_NAME, auto_error=False)
|
|
|
|
|
2022-03-13 04:50:33 +00:00
|
|
|
try:
|
|
|
|
system_key = storage.read_file("system_key.json")["system_key"]
|
|
|
|
except FileNotFoundError:
|
2021-08-23 06:14:40 +00:00
|
|
|
logger.info("System Key not found, generating.")
|
|
|
|
system_uuid = str(uuid.uuid4())
|
2022-03-13 04:50:33 +00:00
|
|
|
storage.write_file("system_key.json", {"system_key": system_uuid})
|
2021-08-23 06:14:40 +00:00
|
|
|
logger.info(f"System UUID: {system_uuid}")
|
|
|
|
|
|
|
|
|
|
|
|
def load_keys():
|
2022-03-13 04:50:33 +00:00
|
|
|
return storage.read_file("api_keys.json")
|
2021-08-23 06:14:40 +00:00
|
|
|
|
|
|
|
|
|
|
|
def get_system_key(
|
|
|
|
api_key_query: str = Security(api_key_query),
|
|
|
|
api_key_header: str = Security(api_key_header),
|
|
|
|
api_key_cookie: str = Security(api_key_cookie)
|
|
|
|
):
|
2022-03-13 04:50:33 +00:00
|
|
|
system_key
|
|
|
|
if api_key_query == system_key:
|
2021-08-23 06:14:40 +00:00
|
|
|
return api_key_query
|
2022-03-13 04:50:33 +00:00
|
|
|
elif api_key_header == system_key:
|
2021-08-23 06:14:40 +00:00
|
|
|
return api_key_header
|
2022-03-13 04:50:33 +00:00
|
|
|
elif api_key_cookie == system_key:
|
2021-08-23 06:14:40 +00:00
|
|
|
return api_key_cookie
|
|
|
|
else:
|
|
|
|
raise HTTPException(
|
|
|
|
status_code=HTTP_403_FORBIDDEN, detail="Could not validate credentials"
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
def get_api_key(
|
|
|
|
api_key_query: str = Security(api_key_query),
|
|
|
|
api_key_header: str = Security(api_key_header),
|
|
|
|
api_key_cookie: str = Security(api_key_cookie)
|
|
|
|
):
|
|
|
|
if api_key_query in load_keys():
|
|
|
|
return api_key_query
|
|
|
|
elif api_key_header in load_keys():
|
|
|
|
return api_key_header
|
|
|
|
elif api_key_cookie in load_keys():
|
|
|
|
return api_key_header
|
|
|
|
else:
|
|
|
|
raise HTTPException(
|
|
|
|
status_code=HTTP_403_FORBIDDEN, detail="Could not validate credentials"
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
@router.get("/key/create")
|
|
|
|
def create_key(name, system_key: APIKey = Depends(get_system_key)):
|
2022-03-13 04:50:33 +00:00
|
|
|
keys: dict = storage.read_file("api_keys.json")
|
2021-08-23 06:14:40 +00:00
|
|
|
new_key = str(uuid.uuid4())
|
|
|
|
keys[new_key] = {}
|
|
|
|
keys[new_key]["name"] = name
|
2022-03-13 04:50:33 +00:00
|
|
|
storage.write_file("api_keys.json", keys)
|
2021-08-23 06:14:40 +00:00
|
|
|
return new_key
|
|
|
|
|
|
|
|
|
|
|
|
@router.get("/key/me")
|
2021-12-01 06:21:26 +00:00
|
|
|
def key_info(access_token: APIKey = Depends(get_api_key)):
|
2021-08-23 06:14:40 +00:00
|
|
|
data = load_keys()
|
2021-12-01 06:21:26 +00:00
|
|
|
return data[access_token]
|