From b977322a4e4cce496594ca183a5fc38a5a31ffc0 Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Tue, 3 May 2016 23:21:01 -0700
Subject: [PATCH] Initial attempt at adding kboot header

---
 atrcopy/__init__.py |  2 +-
 atrcopy/kboot.py    | 22 ++++++++++++++++++++++
 test/test_kboot.py  | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 test/test_kboot.py

diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py
index 1528c1e..018e191 100644
--- a/atrcopy/__init__.py
+++ b/atrcopy/__init__.py
@@ -10,7 +10,7 @@ except ImportError:
 from errors import *
 from ataridos import AtariDosDiskImage, AtariDosFile
 from diskimages import AtrHeader, BootDiskImage
-from kboot import KBootImage
+from kboot import KBootImage, add_kboot_header
 from segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, user_bit_mask, match_bit_mask, comment_bit_mask, data_bit_mask, selected_bit_mask, diff_bit_mask, not_user_bit_mask
 from spartados import SpartaDosDiskImage
 from utils import to_numpy
diff --git a/atrcopy/kboot.py b/atrcopy/kboot.py
index e8a9ba9..4880380 100644
--- a/atrcopy/kboot.py
+++ b/atrcopy/kboot.py
@@ -54,3 +54,25 @@ class KBootImage(DiskImageBase):
         end = dirent.exe_start + dirent.exe_size
         raw = self.rawdata[start:end]
         return XexSegment(raw, 0, 0, start, end, name="KBoot Executable")
+
+kboot_header = '\x96\x020\x13\x80\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x03\x00\x07\x14\x07L\x14\x07\xff\x97\x00\x00\xa9F\x8d\xc6\x02\xd0\xfe\xa0\x00\xa9k\x91X \xd9\x07\xb0\xee \xc4\x07\xadz\x08\rv\x08\xd0\xe3\xa5\x80\x8d\xe0\x02\xa5\x81\x8d\xe1\x02\xa9\x00\x8d\xe2\x02\x8d\xe3\x02 \xeb\x07\xb0\xcc\xa0\x00\x91\x80\xa5\x80\xc5\x82\xd0\x06\xa5\x81\xc5\x83\xf0\x08\xe6\x80\xd0\x02\xe6\x81\xd0\xe3\xadv\x08\xd0\xaf\xad\xe2\x02\x8dp\x07\r\xe3\x02\xf0\x0e\xad\xe3\x02\x8dq\x07 \xff\xff\xadz\x08\xd0\x13\xa9\x00\x8d\xe2\x02\x8d\xe3\x02 \xae\x07\xadz\x08\xd0\x03L<\x07\xa9\x00\x85\x80\x85\x81\x85\x82\x85\x83\xad\xe0\x02\x85\n\x85\x0c\xad\xe1\x02\x85\x0b\x85\r\xa9\x01\x85\t\xa9\x00\x8dD\x02l\xe0\x02 \xeb\x07\x85\x80 \xeb\x07\x85\x81\xa5\x80\xc9\xff\xd0\x10\xa5\x81\xc9\xff\xd0\n \xeb\x07\x85\x80 \xeb\x07\x85\x81 \xeb\x07\x85\x82 \xeb\x07\x85\x83` \xeb\x07\xc9\xff\xd0\t \xeb\x07\xc9\xff\xd0\x02\x18`8`\xad\t\x07\r\n\x07\r\x0b\x07\xf0y\xacy\x08\x10P\xeew\x08\xd0\x03\xeex\x08\xa91\x8d\x00\x03\xa9\x01\x8d\x01\x03\xa9R\x8d\x02\x03\xa9@\x8d\x03\x03\xa9\x80\x8d\x04\x03\xa9\x08\x8d\x05\x03\xa9\x1f\x8d\x06\x03\xa9\x80\x8d\x08\x03\xa9\x00\x8d\t\x03\xadw\x08\x8d\n\x03\xadx\x08\x8d\x0b\x03 Y\xe4\xad\x03\x03\xc9\x02\xb0"\xa0\x00\x8cy\x08\xb9\x80\x08\xaa\xad\t\x07\xd0\x0b\xad\n\x07\xd0\x03\xce\x0b\x07\xce\n\x07\xce\t\x07\xeey\x08\x8a\x18`\xa0\x01\x8cv\x088`\xa0\x01\x8cz\x088`\x00\x03\x00\x80\x00\x00\x00\x00\x00\x00'
+
+def add_kboot_header(bytes):
+    sec_size = 128
+    size = len(bytes)
+    num_sectors = (size + sec_size - 1) / sec_size
+    padded_size = num_sectors * sec_size
+    if size < padded_size:
+        bytes = np.append(bytes, np.zeros([padded_size - size], dtype=np.uint8))
+    paragraphs = padded_size / 16
+    print size, num_sectors, paragraphs
+    header = np.fromstring(kboot_header, dtype=np.uint8)
+    image = np.append(header, bytes)
+    print image
+    words = image.view(dtype='<u2')
+    print words
+    words[1] = paragraphs
+    print words
+    words = image[16 + 9:16 + 9 + 2].view('<u2')
+    words[0] = size
+    return image
diff --git a/test/test_kboot.py b/test/test_kboot.py
new file mode 100644
index 0000000..b4299ab
--- /dev/null
+++ b/test/test_kboot.py
@@ -0,0 +1,35 @@
+import os
+
+import numpy as np
+
+from atrcopy import KBootImage, add_kboot_header
+
+
+class TestKbootHeader(object):
+    def setup(self):
+        pass
+
+    def check_size(self, data):
+        size = np.alen(data)
+        image = add_kboot_header(data)
+        assert int(image[2:4].view(dtype='<u2')) == ((size + 127) / 128) * 128 / 16
+        assert int(image[16 + 9:16 + 9 + 2].view('<u2')) == size
+        return image
+
+    def test_simple(self):
+        for size in range(2000, 40000, 1000):
+            data = np.arange(size, dtype=np.uint8)
+            self.check_size(data)
+    
+    def test_real(self):
+        data = np.fromfile("air_defense_v18.xex", dtype=np.uint8)
+        image = self.check_size(data)
+        with open("air_defense_v18.atr", "wb") as fh:
+            txt = image.tostring()
+            fh.write(txt)
+
+if __name__ == "__main__":
+    t = TestKbootHeader()
+    t.setup()
+    t.test_simple()
+    t.test_real()