1 files changed, 46 insertions, 0 deletions
diff --git a/modules/sts.py b/modules/sts.py
new file mode 100644
index 00000000..07ad2738
--- /dev/null
+++ b/modules/sts.py
@@ -0,0 +1,46 @@
+import time
+from src import ModuleManager, utils
+
+class Module(ModuleManager.BaseModule):
+    def _set_policy(self, server, port, duration):
+        expiration = None
+        if duration:
+            expiration = time.time()+int(duration)
+        server.set_setting("sts-policy", {
+            "port": port,
+            "expiration": expiration})
+    def _change_duration(self, server, info):
+        port = event["server"].port
+        if "port" in info:
+            port = int(info["port"])
+        self._set_policy(server, port, duration)
+
+    @utils.hook("received.cap.ls")
+    def on_cap_ls(self, event):
+        has_sts = "sts" in event["capabilities"]
+        if "sts" in event["capabilities"]:
+            info = utils.parse.keyvalue(event["capabilities"]["sts"],
+                delimiter=",")
+            if not event["server"].tls:
+                self._set_policy(event["server"], int(info["port"]),
+                    None)
+                event["server"].disconnect()
+            else:
+                self._change_duration(event["server"], info)
+
+    @utils.hook("received.cap.new")
+    def on_cap_new(self, event):
+        if "sts" in event["capabilities"] and event["server"].tls:
+            if event["server"].tls:
+                self._change_duration(event["server"], info)
+
+    @utils.hook("new.server")
+    def new_server(self, event):
+        sts_policy = event["server"].get_setting("sts-policy")
+        if sts_policy and not event["server"].tls:
+            expiration = sts_policy["expiration"]
+            if not expiration or time.time() <= expiration:
+                self.log.debug("Applying STS policy for '%s'",
+                    [str(event["server"])])
+                event["server"].tls = True
+                event["server"].port = sts_policy["port"]