120 lines
		
	
	
		
			3.8 KiB
		
	
	
	
		
			Python
		
	
	
	
			
		
		
	
	
			120 lines
		
	
	
		
			3.8 KiB
		
	
	
	
		
			Python
		
	
	
	
| import re
 | |
| import os
 | |
| 
 | |
| 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.
 | |
| 
 | |
|     This function reads specific environment variables and uses their values
 | |
|     to set corresponding security headers. The headers that can be set are:
 | |
|     - cache-control
 | |
|     - strict-transport-security
 | |
|     - referrer-policy
 | |
|     - x-content-type-options
 | |
|     - x-download-options
 | |
|     - x-frame-options
 | |
|     - x-permitted-cross-domain-policies
 | |
| 
 | |
|     Each environment variable is associated with a specific setter function
 | |
|     that constructs the header. If the environment variable is set, the
 | |
|     corresponding header is added to the options dictionary.
 | |
| 
 | |
|     Returns:
 | |
|         dict: A dictionary containing the security headers and their values.
 | |
|     """
 | |
|     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,
 | |
|     }
 | |
| 
 | |
|     for env_var, setter in header_setters.items():
 | |
|         value = os.environ.get(env_var, None)
 | |
|         if value:
 | |
|             header = setter(value)
 | |
|             if header:
 | |
|                 options.update(header)
 | |
| 
 | |
|     return options
 | |
| 
 | |
| # Set HTTP Strict Transport Security(HSTS) response header
 | |
| def set_hsts(value: str):
 | |
|     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
 | |
|     }
 | |
| 
 | |
| # Set X-Frame-Options response header
 | |
| def set_xframe(value: str):
 | |
|     pattern = r'^(DENY|SAMEORIGIN)$'
 | |
|     match = re.match(pattern, value, re.IGNORECASE)
 | |
|     if not match:
 | |
|         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)$'
 | |
|     match = re.match(pattern, value, re.IGNORECASE)
 | |
|     if not match:
 | |
|         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))*$'
 | |
|     match = re.match(pattern, value, re.IGNORECASE)
 | |
|     if not match:
 | |
|         value = 'no-store, max-age=0'
 | |
| 
 | |
|     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
 | |
|     }
 | |
| 
 | |
| # Set X-Content-Type-Options response header
 | |
| def set_xcontent_type(value: str):
 | |
|     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)$'
 | |
|     match = re.match(pattern, value, re.IGNORECASE)
 | |
|     if not match:
 | |
|         value = 'none'
 | |
|     return {
 | |
|         'X-Permitted-Cross-Domain-Policies': value
 | |
|     } |