From 1053863175fa2b770fcc95620c34a4e37fd8ace2 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Thu, 19 Sep 2024 03:24:39 +0200
Subject: [PATCH] chore: format

---
 backend/open_webui/utils/security_headers.py | 79 +++++++++-----------
 1 file changed, 37 insertions(+), 42 deletions(-)

diff --git a/backend/open_webui/utils/security_headers.py b/backend/open_webui/utils/security_headers.py
index c07053917..69a464814 100644
--- a/backend/open_webui/utils/security_headers.py
+++ b/backend/open_webui/utils/security_headers.py
@@ -5,12 +5,14 @@ from fastapi import Request
 from starlette.middleware.base import BaseHTTPMiddleware
 from typing import Dict
 
+
 class SecurityHeadersMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
         response = await call_next(request)
         response.headers.update(set_security_headers())
         return response
 
+
 def set_security_headers() -> Dict[str, str]:
     """
     Sets security headers based on environment variables.
@@ -34,13 +36,13 @@ def set_security_headers() -> Dict[str, str]:
     """
     options = {}
     header_setters = {
-        'CACHE_CONTROL': set_cache_control,
-        'HSTS': set_hsts,
-        'REFERRER_POLICY': set_referrer,
-        'XCONTENT_TYPE': set_xcontent_type,
-        'XDOWNLOAD_OPTIONS': set_xdownload_options,
-        'XFRAME_OPTIONS': set_xframe,
-        'XPERMITTED_CROSS_DOMAIN_POLICIES': set_xpermitted_cross_domain_policies,
+        "CACHE_CONTROL": set_cache_control,
+        "HSTS": set_hsts,
+        "REFERRER_POLICY": set_referrer,
+        "XCONTENT_TYPE": set_xcontent_type,
+        "XDOWNLOAD_OPTIONS": set_xdownload_options,
+        "XFRAME_OPTIONS": set_xframe,
+        "XPERMITTED_CROSS_DOMAIN_POLICIES": set_xpermitted_cross_domain_policies,
     }
 
     for env_var, setter in header_setters.items():
@@ -52,69 +54,62 @@ def set_security_headers() -> Dict[str, str]:
 
     return options
 
+
 # Set HTTP Strict Transport Security(HSTS) response header
 def set_hsts(value: str):
-    pattern = r'^max-age=(\d+)(;includeSubDomains)?(;preload)?$'
+    pattern = r"^max-age=(\d+)(;includeSubDomains)?(;preload)?$"
     match = re.match(pattern, value, re.IGNORECASE)
     if not match:
-        return 'max-age=31536000;includeSubDomains'
-    return {
-        'Strict-Transport-Security': value
-    }
+        return "max-age=31536000;includeSubDomains"
+    return {"Strict-Transport-Security": value}
+
 
 # Set X-Frame-Options response header
 def set_xframe(value: str):
-    pattern = r'^(DENY|SAMEORIGIN)$'
+    pattern = r"^(DENY|SAMEORIGIN)$"
     match = re.match(pattern, value, re.IGNORECASE)
     if not match:
-        value = 'DENY'
-    return {
-        "X-Frame-Options": value
-    }
+        value = "DENY"
+    return {"X-Frame-Options": value}
+
 
 # Set Referrer-Policy response header
 def set_referrer(value: str):
-    pattern = r'^(no-referrer|no-referrer-when-downgrade|origin|origin-when-cross-origin|same-origin|strict-origin|strict-origin-when-cross-origin|unsafe-url)$'
+    pattern = r"^(no-referrer|no-referrer-when-downgrade|origin|origin-when-cross-origin|same-origin|strict-origin|strict-origin-when-cross-origin|unsafe-url)$"
     match = re.match(pattern, value, re.IGNORECASE)
     if not match:
-        value = 'no-referrer'
-    return {
-        'Referrer-Policy': value
-    }
+        value = "no-referrer"
+    return {"Referrer-Policy": value}
+
 
 # Set Cache-Control response header
 def set_cache_control(value: str):
-    pattern = r'^(public|private|no-cache|no-store|must-revalidate|proxy-revalidate|max-age=\d+|s-maxage=\d+|no-transform|immutable)(,\s*(public|private|no-cache|no-store|must-revalidate|proxy-revalidate|max-age=\d+|s-maxage=\d+|no-transform|immutable))*$'
+    pattern = r"^(public|private|no-cache|no-store|must-revalidate|proxy-revalidate|max-age=\d+|s-maxage=\d+|no-transform|immutable)(,\s*(public|private|no-cache|no-store|must-revalidate|proxy-revalidate|max-age=\d+|s-maxage=\d+|no-transform|immutable))*$"
     match = re.match(pattern, value, re.IGNORECASE)
     if not match:
-        value = 'no-store, max-age=0'
+        value = "no-store, max-age=0"
+
+    return {"Cache-Control": value}
 
-    return {
-        'Cache-Control': value
-    }
 
 # Set X-Download-Options response header
 def set_xdownload_options(value: str):
-    if value != 'noopen':
-        value = 'noopen'
-    return {
-        'X-Download-Options': value
-    }
+    if value != "noopen":
+        value = "noopen"
+    return {"X-Download-Options": value}
+
 
 # Set X-Content-Type-Options response header
 def set_xcontent_type(value: str):
-    if value != 'nosniff':
-        value = 'nosniff'
-    return {
-        'X-Content-Type-Options': value
-    }
+    if value != "nosniff":
+        value = "nosniff"
+    return {"X-Content-Type-Options": value}
+
 
 # Set X-Permitted-Cross-Domain-Policies response header
 def set_xpermitted_cross_domain_policies(value: str):
-    pattern = r'^(none|master-only|by-content-type|by-ftp-filename)$'
+    pattern = r"^(none|master-only|by-content-type|by-ftp-filename)$"
     match = re.match(pattern, value, re.IGNORECASE)
     if not match:
-        value = 'none'
-    return {
-        'X-Permitted-Cross-Domain-Policies': value
-    }
\ No newline at end of file
+        value = "none"
+    return {"X-Permitted-Cross-Domain-Policies": value}