From 8ebad01de8cbbec82323a0c1a78366a30b803b65 Mon Sep 17 00:00:00 2001
From: Connor <cloningdonor@gmail.com>
Date: Wed, 30 Aug 2017 21:44:22 -0700
Subject: [PATCH] .

---
 bwmap.py | 40 ++++++++++++++--------------------------
 1 file changed, 14 insertions(+), 26 deletions(-)

diff --git a/bwmap.py b/bwmap.py
index 7792520..3db4e6b 100644
--- a/bwmap.py
+++ b/bwmap.py
@@ -21,8 +21,10 @@ scmdraft_req = b"ERA ~DIM ~TILE~UNIT~PUNI~IOWN~THG2~SPRP~WAV ~MASK~STR ~OWNR~SID
 # STR  kinda (can be empty for melee)
 melee_req = b"VER ~TYPE~VCOD~OWNR~SIDE~COLR~ERA ~DIM ~MTXM~UNIT~THG2~STR ~SPRP~FORC".split(b"~")
 
-#both_req = set(melee_req + scmdraft_req)
-#print(b"~".join(both_req))
+# TODO: try more things to add to this list.
+ums_req = melee_req + b"MRGN~PTEx~PUPx~TECx~TRIG~UNIx~UPGx~UPRP".split(b"~")
+
+both_req = set(melee_req + scmdraft_req)
 
 known_keys = [ # in the order we want to write them.
     # boring filetype stuff
@@ -79,15 +81,7 @@ known_keys = [ # in the order we want to write them.
     b"STR ", # string data
 ]
 
-skippable_keys = [
-    b"ISOM",
-    b"PTEC",
-    b"SWNM",
-    b"TECS",
-    b"UNIS",
-    b"UPGR",
-    b"UPGS",
-]
+skippable_keys = set(known_keys).symmetric_difference(both_req)
 
 vcod_valid = b"\
 \x34\x19\xCA\x77\x99\xDC\x68\x71\x0A\x60\xBF\xC3\xA7\xE7\x75\xA7\
@@ -175,7 +169,7 @@ def readit(f):
         k = f.read(4)
 
         try:
-            name = k.decode('utf-8')#, errors="replace")
+            name = k.decode('ascii')
         except UnicodeDecodeError:
             print("! skipping invalid section name:", k)
             size = U("<l", f.read(4))[0]
@@ -194,7 +188,6 @@ def readit(f):
             print("expected:", size)
             return
 
-        #print("heresay: {:X}".format(f.tell()))
         #if k not in b"ERA ~DIM ~MTXM~THG2~UNIT".split(b"~"): continue
 
         if k not in src:
@@ -215,7 +208,6 @@ def makestringtable(strings):
     o = count * 2 + 3
     for s in strings:
         offsets += [o]
-        # FIXME: offsets are still broken.
         o += len(s) + 1
 
     out = P("<H", count)
@@ -229,8 +221,7 @@ def rewrite(src):
     if b"STR " not in src:
         return
 
-    ok = b"MRGN~MBRF~TRIG~SPRP~FORC~WAV ~UNIx~SWNM"
-    ok = ok.split(b"~")
+    ok = b"MRGN~MBRF~TRIG~SPRP~FORC~WAV ~UNIx~SWNM".split(b"~")
 
     refs = []
 
@@ -414,17 +405,16 @@ def writeit(f, src):
     assert b"DIM " in src # TODO: try to guess map dimensions.
 
     if b"COLR" not in src:
-        #src[b"COLR"] = [b"\0\1\2\3\4\5\6\7\13\13\13\13"]
         src[b"COLR"] = [b"\0\1\2\3\4\5\6\7"]
 
-    if b"MASK" not in src:
-        # SCMDraft 2 seems to require this to exist and be the correct length.
+    # SCMDraft 2 seems to require this to be the correct length.
+    if b"MASK" not in src or len(src[b"MASK"][0]) != w * h:
         src[b"MASK"] = [b"\xFF" * (w * h)]
 
     if b"STR " not in src:
         src[b"STR "] = [b""]
 
-    # SCMDraft 2 calls these both "player settings", so i'm grouping then.
+    # SCMDraft 2 calls this "player settings" in its errors.
     if b"OWNR" not in src:
         b = b""
         for i in range(12):
@@ -432,13 +422,8 @@ def writeit(f, src):
         src[b"OWNR"] = [b]
         assert len(b) == 12, b
 
+    # SCMDraft 2 calls this "player settings" in its errors.
     if b"SIDE" not in src:
-    #   b = b""
-    #   for i in range(8):
-    #       b += b"\5" if placed_starts[i] else b"\6"
-    #   b += b"\7\7\7\4"
-    #   src[b"SIDE"] = [b]
-    #   assert len(b) == 12, b
         src[b"SIDE"] = [b"\5\5\5\5\5\5\5\5\7\7\7\4"]
 
     if b"FORC" not in src:
@@ -500,6 +485,9 @@ def writeit(f, src):
     for k in melee_req:
         if k not in wroteonce:
             print("# map might not open in StarCraft (Melee). missing", k)
+    for k in ums_req:
+        if k not in wroteonce:
+            print("# map might not open in StarCraft (Use Map Settings). missing", k)
 
 def run(args):
     for fn in args: