From 3b7d46e5a726de5265bafcba8b6b81069966e805 Mon Sep 17 00:00:00 2001
From: ChillerDragon <ChillerDragon@gmail.com>
Date: Tue, 9 May 2023 10:04:37 +0200
Subject: [PATCH] First working version of map download example

---
 examples/download_map/.gitignore      |  1 +
 examples/download_map/download_map.py | 53 +++++++++++++++++++++++----
 2 files changed, 46 insertions(+), 8 deletions(-)
 create mode 100644 examples/download_map/.gitignore

diff --git a/examples/download_map/.gitignore b/examples/download_map/.gitignore
new file mode 100644
index 0000000..c7de605
--- /dev/null
+++ b/examples/download_map/.gitignore
@@ -0,0 +1 @@
+*.map
diff --git a/examples/download_map/download_map.py b/examples/download_map/download_map.py
index 2c34df3..0733a85 100755
--- a/examples/download_map/download_map.py
+++ b/examples/download_map/download_map.py
@@ -1,12 +1,15 @@
 #!/usr/bin/env python
 
-import socket, os
-
+import socket
+import os
+import sys
+import hashlib
 from typing import Optional
 
 from twnet_parser.packet import TwPacket, parse7
 from twnet_parser.messages7.control.token import CtrlToken
 from twnet_parser.messages7.control.connect import CtrlConnect
+from twnet_parser.messages7.control.close import CtrlClose
 from twnet_parser.messages7.system.info import MsgInfo
 from twnet_parser.messages7.system.map_change import MsgMapChange
 from twnet_parser.messages7.system.request_map_data import MsgRequestMapData
@@ -22,14 +25,28 @@ def progress(have, want):
     bar = '=' * int(percent / 2)
     bar = bar[0:-1] + '>'
     print(f"[{bar:50}] {have}/{want} {percent}%")
+    print(outfile)
 
-host = 'localhost'
-port = 8303
+if len(sys.argv) < 3:
+    print("usage: download.py HOST PORT")
+    print("description:")
+    print("  connects to given server and downloads the map")
+    print("  the file will be stored in the current working directory")
+    print("  then the client disconnects and quits the program")
+    print("example:")
+    print("  download.py localhost 8303")
+    exit(1)
+
+host = sys.argv[1]
+port = int(sys.argv[2])
 dest_srv = (host, port)
 
+print(f"Connecting to {host}:{port} ...")
+
 sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 sock.bind(('', 0))
 
+outfile = None
 downloaded_chunks = 0
 downloaded_bytes = 0
 sent_vital = 0
@@ -62,6 +79,17 @@ ctrl_token = CtrlToken()
 ctrl_token.response_token = my_token
 send_msg(ctrl_token)
 
+def verify_checksum(filename, sha256):
+    with open(filename, "rb") as f:
+        bytes = f.read()
+        readable_hash = hashlib.sha256(bytes).hexdigest()
+        if readable_hash == sha256:
+            print("OK checksum matches.")
+            return
+        else:
+            print(f"ERROR expected={sha256} got={readable_hash}")
+    exit(1)
+
 while True:
     data, addr = sock.recvfrom(NET_MAX_PACKETSIZE)
     packet = parse7(data)
@@ -71,8 +99,7 @@ while True:
             if msg.header.flags.vital and not msg.header.flags.resend:
                 got_vital = (got_vital + 1) % NET_MAX_SEQUENCE
             if msg.header.seq in got_seqs:
-                pass
-                # continue
+                continue
             got_seqs.add(msg.header.seq)
 
         if msg.message_name == 'token':
@@ -87,20 +114,30 @@ while True:
             send_msg(info)
         elif msg.message_name == 'map_change':
             map_info = msg
+            outfile = f"{map_info.name}_{map_info.sha256.hex()}.map"
+            if os.path.isfile(outfile):
+                print(f"Error: file already exists '{outfile}'")
+                exit(1)
             req = MsgRequestMapData()
             req.header.flags.vital = True
             send_msg(req)
         elif msg.message_name == 'map_data':
+            if not outfile:
+                continue
             downloaded_chunks += 1
             downloaded_bytes += msg.header.size
+            with open(outfile, "ab") as map_file:
+                map_file.write(msg.data)
             progress(downloaded_bytes, map_info.size)
             if downloaded_bytes >= map_info.size:
                 print(f"finished downloading {map_info.name}")
-                continue
+                verify_checksum(outfile, map_info.sha256.hex())
+                send_msg(CtrlClose())
+                exit(0)
             if downloaded_chunks % map_info.num_response_chunks_per_request == 0:
                 req = MsgRequestMapData()
                 req.header.flags.vital = True
                 send_msg(req)
         elif msg.message_name == 'close':
             print("disconnected")
-            exit(0)
+            exit(1)