Added Host Access Control

main.py

@@ -1,36 +1,42 @@
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.middleware.trustedhost import TrustedHostMiddleware
 from QuoteGenerator import QuoteGenerator
 from typing import Union
 from pydantic import BaseModel
+import time
 
 class QuoteRequest(BaseModel):
     tags: Union[None, str]
 
 app = FastAPI()
 
-origins = ['http://testinghost.com']
+origins = ['quote-generator.streamlit.app']
 
-# # Setting up cors
+# Middleware to valiate host
 app.add_middleware(
-    CORSMiddleware,
-    allow_origins=[''],
-    allow_credentials=True,
-    allow_methods=["POST"],
-    allow_headers=["Content-Type"]
+    TrustedHostMiddleware,
+    allowed_hosts=origins
 )
 
+#Middleware to note time
+@app.middleware("http")
+async def note_response_time(request: Request, call_next):
+    start_time = time.time()
+    response = await call_next(request)
+    process_time = time.time()
+    print(f'Time taken = {process_time-start_time:.1f}s')
+    return response
+
 quote_generator = QuoteGenerator()
 quote_generator.load_generator()
 
 @app.get("/")
 def root(request: Request):
-    print("URL ", request.client.host)
     return {"message": "This is the website for quote-generator"}
 
 @app.post("/generate_quote")
-def generate_quote(req: QuoteRequest, reqs: Request):
+def generate_quote(req: QuoteRequest):
     print("Tags: ", req.tags)
-    print("Path: ", reqs.headers.get("Referer"))
-    # generated_quote_oup = quote_generator.generate_quote(req.tags)
-    return {'quote': "Generated quote"}
+    generated_quote_oup = quote_generator.generate_quote(req.tags)
+    return {'quote': generated_quote_oup}