hsts magic

flask_sslify.py

@@ -3,9 +3,9 @@
 from flask import request, redirect
 
 class SSLify(object):
+    """Secures your Flask App."""
 
     def __init__(self, app, age=31536000, subdomains=False):
-
         if app is not None:
             self.app = app
             self.hsts_age = age
@@ -16,10 +16,23 @@ class SSLify(object):
             self.app = None
 
     def init_app(self, app):
-        app.before_request(self.redirect)
+        """Configures the configured Flask app to enforce SSL."""
+        app.before_request(self.redirect_to_ssl)
+        app.after_request(self.set_hsts_header)
 
-    def redirect(self):
+    @property
+    def hsts_header(self):
+        """Returns the proper HSTS policy."""
+        hsts_policy = 'max-age={0}'.format(self.hsts_age)
 
+        if self.hsts_include_subdomains:
+            hsts_policy += '; includeSubDomains'
+
+        return hsts_policy
+
+    def redirect_to_ssl(self):
+        """Redirect incoming requests to HTTPS."""
+        # Should we redirect?
         criteria = [
             request.is_secure,
             self.app.debug,
@@ -30,14 +43,13 @@ class SSLify(object):
             url = request.url.replace('http://', 'https://')
             r = redirect(url)
 
-            # HSTS policy.
-            hsts_policy = 'max-age={0}'.format(self.hsts_age)
-
-            if self.hsts_include_subdomains:
-                hsts_policy += '; includeSubDomains'
-
-            r.headers['Strict-Transport-Security'] = hsts_policy
-
+            r.headers['Strict-Transport-Security'] = self.hsts_header
             return r
 
+    def set_hsts_header(self, response):
+        """Adds HSTS header to each response."""
+        response.headers.setdefault('Strict-Transport-Security', self.hsts_header)
+        return response
+
+