Made STS policies one-shot when upgrading from plaintext to tls

1 files changed, 11 insertions, 4 deletions

diff --git a/modules/sts.py b/modules/sts.py
index 98da6581..30d9678f 100644
--- a/modules/sts.py
+++ b/modules/sts.py
@@ -2,18 +2,19 @@ import time
 from src import ModuleManager, utils
 
 class Module(ModuleManager.BaseModule):
-    def _set_policy(self, server, port, duration):
+    def _set_policy(self, server, port, duration, one_shot):
         expiration = None
         if duration:
             expiration = time.time()+int(duration)
         server.set_setting("sts-policy", {
             "port": port,
-            "expiration": expiration})
+            "expiration": expiration,
+            "one-shot": one_shot})
     def _change_duration(self, server, info):
         port = event["server"].port
         if "port" in info:
             port = int(info["port"])
-        self._set_policy(server, port, info["duration"])
+        self._set_policy(server, port, info["duration"], False)
 
     @utils.hook("received.cap.ls")
     def on_cap_ls(self, event):
@@ -23,7 +24,7 @@ class Module(ModuleManager.BaseModule):
                 delimiter=",")
             if not event["server"].tls:
                 self._set_policy(event["server"], int(info["port"]),
-                    None)
+                    None, True)
                 event["server"].disconnect()
             else:
                 self._change_duration(event["server"], info)
@@ -44,3 +45,9 @@ class Module(ModuleManager.BaseModule):
                     [str(event["server"])])
                 event["server"].tls = True
                 event["server"].port = sts_policy["port"]
+
+    @utils.hook("received.numeric.001"):
+    def on_connect(self, event):
+        sts_policy = event["server"].get_setting("sts-policy")
+        if sts_policy and sts_policy["one-shot"]:
+            event["server"].del_setting("sts-policy")