From 07cdb05ba20575d5b7d44884b9783c647eaed92f Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Sun, 26 Feb 2017 12:14:23 -0800
Subject: [PATCH] Added -s command to use pyatasm to assemble MAC/65 source and
 place it in a file in the disk image

---
 atrcopy/__init__.py   | 68 +++++++++++++++++++++++++++++++++----------
 atrcopy/ataridos.py   | 24 ++++++++++-----
 atrcopy/diskimages.py |  3 ++
 atrcopy/dos33.py      | 22 ++++++++++++++
 4 files changed, 94 insertions(+), 23 deletions(-)

diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py
index c3ec609..ada9333 100644
--- a/atrcopy/__init__.py
+++ b/atrcopy/__init__.py
@@ -90,27 +90,32 @@ def extract_files(image, files):
             with open(name, "wb") as fh:
                 fh.write(data)
 
+def save_file(image, name, filetype, data):
+    try:
+        dirent = image.find_dirent(name)
+        if options.force:
+            image.delete_file(name)
+        else:
+            print "skipping %s, use -f to overwrite" % (name)
+            return False
+    except FileNotFound:
+        pass
+    print "copying %s to %s" % (name, image)
+    if not options.dry_run:
+        image.write_file(name, filetype, data)
+        return True
+    return False
+
+
 def add_files(image, files):
     filetype = options.filetype
     if not filetype:
         filetype = image.default_filetype
     changed = False
     for name in files:
-        try:
-            dirent = image.find_dirent(name)
-            if options.force:
-                image.delete_file(name)
-            else:
-                print "skipping %s, use -f to overwrite" % (name)
-                continue
-        except FileNotFound:
-            pass
-        print "copying %s to %s" % (name, image)
-        if not options.dry_run:
-            with open(name, "rb") as fh:
-                data = fh.read()
-            image.write_file(name, filetype, data)
-            changed = True
+        with open(name, "rb") as fh:
+            data = fh.read()
+        changed = save_file(image, name, filetype, data)
     if changed:
         image.save()
 
@@ -135,6 +140,32 @@ def list_files(image, files):
         if not files or dirent.filename in files:
             print dirent
 
+def assemble(image, files):
+    try:
+        import pyatasm
+    except ImportError:
+        raise AtrError("Please install pyatasm to compile code.")
+    changed = False
+    segments = []
+    print files
+    for name in files:
+        try:
+            asm = pyatasm.Assemble(name)
+        except SyntaxError, e:
+            raise AtrError("Assembly error: %s" % e.msg)
+        for first, last, object_code in asm.segments:
+            print "%04x - %04x, size=%04x" % (first, last, len(object_code))
+            rawdata = SegmentData(object_code)
+            s = DefaultSegment(rawdata, first, name)
+            segments.append(s)
+    if options.verbose:
+        for s in segments:
+            print "%s - %04x)" % (str(s)[:-1], s.start_addr + len(s))
+    file_data, filetype = image.create_executable_file_image(segments)
+    changed = save_file(image, options.output, filetype, file_data)
+    if changed:
+        image.save()
+
 
 def run():
     import sys
@@ -152,12 +183,15 @@ def run():
     parser.add_argument("--xex", action="store_true", default=False, help="add .xex extension")
     parser.add_argument("-f", "--force", action="store_true", default=False, help="force operation, allowing file overwrites and operation on non-standard disk images")
     parser.add_argument("files", metavar="IMAGE", nargs="+", help="a disk image file [or a list of them]")
-    parser.add_argument("-s", "--segments", action="store_true", default=False, help="display segments")
+    parser.add_argument("-g", "--segments", action="store_true", default=False, help="display segments")
     parser.add_argument("-x", "-e", "--extract", action="store_true", default=False, help="extract named files")
     parser.add_argument("-a", "--add", action="store_true", default=False, help="add files to image")
     parser.add_argument("-r", "--remove", action="store_true", default=False, help="remove named files from image")
     parser.add_argument("-t", "--filetype", action="store", default="", help="file type metadata for writing to disk images that require it")
+    parser.add_argument("-s", "--asm", nargs="+", action="append", help="source file(s) to assemble using pyatasm (requires -o to specify filename stored on disk image)")
+    parser.add_argument("-o", "--output", action="store", default="", help="output file name for those commands that need it")
     options, extra_args = parser.parse_known_args()
+    print options, extra_args
 
     # Turn off debug messages by default
     logging.basicConfig(level=logging.WARNING)
@@ -186,5 +220,7 @@ def run():
                 extract_files(parser.image, file_list)
             elif options.remove:
                 remove_files(parser.image, file_list)
+            elif options.asm:
+                assemble(parser.image, options.asm[0])
             else:
                 list_files(parser.image, file_list)
diff --git a/atrcopy/ataridos.py b/atrcopy/ataridos.py
index abd6708..2abbb32 100644
--- a/atrcopy/ataridos.py
+++ b/atrcopy/ataridos.py
@@ -2,7 +2,7 @@ import numpy as np
 
 from errors import *
 from diskimages import DiskImageBase, BaseHeader
-from segments import EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver, get_style_bits
+from segments import SegmentData, EmptySegment, ObjSegment, RawSectorsSegment, DefaultSegment, SegmentSaver, get_style_bits
 from utils import *
 
 import logging
@@ -616,6 +616,9 @@ class AtariDosDiskImage(DiskImageBase):
                 log.debug("%s not a binary file; skipping segment generation" % str(segment))
         return segments_out
 
+    def create_executable_file_image(self, segments,run_addr=None):
+        return get_xex(segments, run_addr), "XEX"
+
 
 class BootDiskImage(AtariDosDiskImage):
     def __str__(self):
@@ -645,11 +648,22 @@ class BootDiskImage(AtariDosDiskImage):
             raise InvalidDiskImage("Bad boot load address")
 
 
-def get_xex(segments, runaddr):
+def get_xex(segments, runaddr=None):
     total = 2
+    runad = False
     for s in segments:
         total += 4 + len(s)
-    total += 6
+        if s.start_addr == 0x2e0:
+            runad = True
+    if not runad:
+        words = np.empty([1], dtype='<u2')
+        if not runaddr:
+            runaddr = segments[0].start_addr
+        words[0] = runaddr
+        r = SegmentData(words.view(dtype=np.uint8))
+        s = DefaultSegment(r, 0x2e0)
+        segments[0:0] = [s]
+        total += 6
     bytes = np.zeros([total], dtype=np.uint8)
     bytes[0:2] = 0xff # FFFF header
     i = 2
@@ -660,10 +674,6 @@ def get_xex(segments, runaddr):
         i += 4
         bytes[i:i + len(s)] = s[:]
         i += len(s)
-    words = bytes[i:i+6].view(dtype='<u2')
-    words[0] = 0x2e0
-    words[1] = 0x2e1
-    words[2] = runaddr
     return bytes
 
 def add_atr_header(bytes):
diff --git a/atrcopy/diskimages.py b/atrcopy/diskimages.py
index d8d349a..33a3234 100644
--- a/atrcopy/diskimages.py
+++ b/atrcopy/diskimages.py
@@ -301,6 +301,9 @@ class DiskImageBase(object):
             segments.append(segment)
         return segments
 
+    def create_executable_file_image(self, segments):
+        raise NotImplementedError
+
     # file writing methods
 
     def begin_transaction(self):
diff --git a/atrcopy/dos33.py b/atrcopy/dos33.py
index f7a1fe5..f1adb18 100644
--- a/atrcopy/dos33.py
+++ b/atrcopy/dos33.py
@@ -537,6 +537,28 @@ class Dos33DiskImage(DiskImageBase):
             segment = EmptySegment(self.rawdata, name=dirent.filename)
         return segment
 
+    def create_executable_file_image(self, segments):
+        origin = 100000000
+        last = -1
+        for s in segments:
+            origin = min(origin, s.start_addr)
+            last = max(last, s.start_addr + len(s))
+            print "contiguous bytes needed: %04x - %04x" % (origin, last)
+        size = last - origin
+        image = np.zeros([size + 4], dtype=np.uint8)
+        words = image[0:4].view(dtype="<u2")  # always little endian
+        words[0] = origin
+        words[1] = size
+        for s in segments:
+            index = s.start_addr - origin + 4
+            print "setting data for %04x - %04x at %04x" % (s.start_addr, s.start_addr + len(s), index)
+            image[index:index + len(s)] = s.data
+        return image, 'B'
+
+
+
+
+
 
 class ProdosHeader(Dos33Header):
     file_format = "ProDOS"