Skip to content

Commit 1367d95

Browse files
committed
refac: enforce api on all routes
1 parent c98ca76 commit 1367d95

File tree

2 files changed

+23
-8
lines changed

2 files changed

+23
-8
lines changed

main.py

+11-8
Original file line numberDiff line numberDiff line change
@@ -106,28 +106,31 @@ def get_all_pipelines():
106106

107107
return pipelines
108108

109+
109110
def parse_frontmatter(content):
110111
frontmatter = {}
111-
for line in content.split('\n'):
112-
if ':' in line:
113-
key, value = line.split(':', 1)
112+
for line in content.split("\n"):
113+
if ":" in line:
114+
key, value = line.split(":", 1)
114115
frontmatter[key.strip().lower()] = value.strip()
115116
return frontmatter
116117

118+
117119
def install_frontmatter_requirements(requirements):
118120
if requirements:
119-
req_list = [req.strip() for req in requirements.split(',')]
121+
req_list = [req.strip() for req in requirements.split(",")]
120122
for req in req_list:
121123
print(f"Installing requirement: {req}")
122124
subprocess.check_call([sys.executable, "-m", "pip", "install", req])
123125
else:
124126
print("No requirements found in frontmatter.")
125127

128+
126129
async def load_module_from_path(module_name, module_path):
127130

128131
try:
129132
# Read the module content
130-
with open(module_path, 'r') as file:
133+
with open(module_path, "r") as file:
131134
content = file.read()
132135

133136
# Parse frontmatter
@@ -139,8 +142,8 @@ async def load_module_from_path(module_name, module_path):
139142
frontmatter = parse_frontmatter(frontmatter_content)
140143

141144
# Install requirements if specified
142-
if 'requirements' in frontmatter:
143-
install_frontmatter_requirements(frontmatter['requirements'])
145+
if "requirements" in frontmatter:
146+
install_frontmatter_requirements(frontmatter["requirements"])
144147

145148
# Load the module
146149
spec = importlib.util.spec_from_file_location(module_name, module_path)
@@ -277,7 +280,7 @@ async def check_url(request: Request, call_next):
277280

278281
@app.get("/v1/models")
279282
@app.get("/models")
280-
async def get_models():
283+
async def get_models(user: str = Depends(get_current_user)):
281284
"""
282285
Returns the available pipelines
283286
"""

utils/pipelines/auth.py

+12
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
22
from fastapi import HTTPException, status, Depends
33

4+
45
from pydantic import BaseModel
56
from typing import Union, Optional
67

@@ -14,6 +15,10 @@
1415
import requests
1516
import uuid
1617

18+
19+
from config import API_KEY, PIPELINES_DIR
20+
21+
1722
SESSION_SECRET = os.getenv("SESSION_SECRET", " ")
1823
ALGORITHM = "HS256"
1924

@@ -62,4 +67,11 @@ def get_current_user(
6267
credentials: HTTPAuthorizationCredentials = Depends(bearer_security),
6368
) -> Optional[dict]:
6469
token = credentials.credentials
70+
71+
if token != API_KEY:
72+
raise HTTPException(
73+
status_code=status.HTTP_401_UNAUTHORIZED,
74+
detail="Invalid API key",
75+
)
76+
6577
return token

0 commit comments

Comments
 (0)