include timestamp into message and signature

Snawoot · Aug 13, 2017 · 8bc9a28 · 8bc9a28
1 parent 177e0a7
commit 8bc9a28
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 16 deletions.
diff --git a/pyknock.py b/pyknock.py
@@ -5,6 +5,11 @@
 import hashlib
 import hmac
 import argparse
+import time
+import struct
+
+CODE_OPEN = 1
+CODE_CLOSE = 2
 
 def main():
     parser = argparse.ArgumentParser()
@@ -25,9 +30,11 @@ def main():
         myip = socket.inet_pton(socket.AF_INET, args.sign_address)
     else:
         myip = socket.inet_pton(socket.AF_INET, s.getsockname()[0])
-    opcode = '\x01' if args.command == 'open' else '\x02'
-    digest = hmac.new(args.psk, opcode + myip, hashlib.sha256).digest()
-    s.sendall(opcode + digest)
+
+    opcode = CODE_OPEN if args.command == 'open' else CODE_CLOSE
+    msg = struct.pack('<Bd', opcode, time.time())
+    digest = hmac.new(args.psk, msg + myip, hashlib.sha256).digest()
+    s.sendall(msg + digest)
 
 if __name__ == '__main__':
     main()
diff --git a/pyknockd.py b/pyknockd.py
@@ -7,11 +7,22 @@
 import hmac
 import string
 import argparse
+import struct
+import time
+
+DIGEST = hashlib.sha256
+DIGEST_SIZE = DIGEST().digest_size
+BODY_SIZE = 1 + 8
+MSG_SIZE = BODY_SIZE + DIGEST_SIZE
+CODE_OPEN = 1
+CODE_CLOSE = 2
+MSG_FMT = "<Bd%ds" % (DIGEST_SIZE)
 
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("-b", "--bind-address", help="bind address", default="")
     parser.add_argument("-p", "--port", help="bind port", type=int, default=60120, metavar="PORT")
+    parser.add_argument("-t", "--time-drift", help="allowed time drift in seconds between client and server. Value may be a floating point number", type=float, default=60, metavar="DRIFT")
     parser.add_argument("psk", help="pre-shared key used to authenticate clients", metavar="PSK")
     parser.add_argument("open_cmd", help="template of command used to enable access. Example: \"ipset add -exist myset $ip\". Available variables: $ip, $port", metavar="OPEN_CMD")
     parser.add_argument("close_cmd", help="template of command used to disable access. Example: \"ipset del -exist myset $ip\". Available variables: $ip, $port", metavar="CLOSE_CMD")
@@ -25,24 +36,33 @@ def main():
         while True:
             data, addr = s.recvfrom(4096)
             ip4 = socket.inet_pton(socket.AF_INET, addr[0])
+
             if not data:
                 continue
 
-            opcode = data[0]
-            signature = data[1:]
-            if hmac.compare_digest(
+            if len(data) != MSG_SIZE:
+                continue
+
+            opcode, ts, digest = struct.unpack(MSG_FMT, data)
+
+            if abs(ts - time.time()) > args.time_drift:
+                continue
+
+            if not hmac.compare_digest(
                     hmac.new(
-                            args.psk, opcode + ip4, hashlib.sha256
+                            args.psk, data[:BODY_SIZE] + ip4, DIGEST
                         ).digest(),
-                    signature):
-                if opcode == '\x01':
-                    os.system(
-                        open_cmd.safe_substitute(ip = addr[0], port = addr[1])
-                    )
-                elif opcode == '\x02':
-                    os.system(
-                        close_cmd.safe_substitute(ip = addr[0], port = addr[1])
-                    )
+                    digest):
+                continue
+
+            if opcode == CODE_OPEN:
+                os.system(
+                    open_cmd.safe_substitute(ip = addr[0], port = addr[1])
+                )
+            elif opcode == CODE_CLOSE:
+                os.system(
+                    close_cmd.safe_substitute(ip = addr[0], port = addr[1])
+                )
     except KeyboardInterrupt:
         pass