merge soas-assimilator.py and gtktest.py ; some minor fixes

2 files changed, 217 insertions, 270 deletions

diff --git a/gtktest.py b/gtktest.py
deleted file mode 100755
index 5bb919a..0000000
--- a/gtktest.py
+++ /dev/null
@@ -1,215 +0,0 @@
-#!/usr/bin/env python
-import os
-import sys
-import gobject
-import gtk
-
-PROGNAME = "soas-assimilator"
-
-
-class tColoredListWindow (gtk.Window) :
-
-  def __init__(self, title) :
-    gtk.Window.__init__(self)
-    self.set_title(title)
-    self.initWidgets()
-
-  def initWidgets(self) :
-    self.vbox = gtk.VBox(False, 10)
-    self.add(self.vbox)
-    self._scrolledList = gtk.ScrolledWindow()
-    self._scrolledList.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
-    self.vbox.pack_start(self._scrolledList)
-    self.listModel = gtk.ListStore(*[_type for (name, _type, r) in self.listColumns])
-    self._listView = gtk.TreeView(self.listModel)
-    self._scrolledList.add(self._listView)
-    self._addListColumns()
-    self.show_all()
-
-  def _addListColumns(self) :
-    for (idx, (name, _type, renderer)) in enumerate(self.listColumns) :
-      if not renderer :
-        continue
-
-      r = renderer()
-      r.set_property('background-set', True)
-      if (_type == bool) :
-        col = gtk.TreeViewColumn(name, r, active=idx, background=self.colorColIdx)
-      else :
-        col = gtk.TreeViewColumn(name, r, text=idx, background=self.colorColIdx)
-
-      col.set_sort_column_id(idx)
-      self._listView.append_column(col)
-
-  def findRow(self, col, content) :
-    for idx in range(len(self.listModel)) :
-      if (self.listModel[idx][col] == content) :
-        return idx
-
-
-class tReadWindow (tColoredListWindow) :
-
-  listColumns = [
-    ('status', str, gtk.CellRendererText),
-    ('image', str, gtk.CellRendererText),
-    ('model', str, gtk.CellRendererText),
-    ('device', str, None),
-    ('color', str, None),
-  ]
-  _deviceColIdx = 3
-  colorColIdx = len(listColumns)-1
-  _colorMap = {
-    'reading': '#FFFF00',
-    'finished': '#00FF00',
-    'error': '#FF0000',
-  }
-
-  def __init__(self, model) :
-    self._model = model
-    tColoredListWindow.__init__(self, "%s: Read USB sticks" % (PROGNAME,))
-
-  def initWidgets(self) :
-    tColoredListWindow.initWidgets(self)
-    self._hButtonBox = gtk.HButtonBox()
-    self.writeButton = gtk.Button("Continue to write mode")
-    self.writeButton.modify_bg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#F00000"))
-    self.writeButton.modify_bg(gtk.STATE_PRELIGHT, gtk.gdk.color_parse("#FF0000"))
-    self.writeButton.show()
-    self._hButtonBox.pack_start(self.writeButton)
-    self._hButtonBox.show()
-    self.vbox.pack_end(self._hButtonBox, expand=False)
-
-  def setStatus(self, device, status, image, model) :
-    row = self.findRow(device)
-    if row is None :
-      row = self.listModel.append()
-
-    color = self._colorMap[status]
-    self.listModel[row] = [locals()[name] for (name, t, r) in self.listColumns]
-
-  def findRow(self, device) :
-    return tColoredListWindow.findRow(self, self._deviceColIdx, device)
-
-
-class tWriteWindow (tColoredListWindow) :
-
-  listColumns = [
-    ('device', str, gtk.CellRendererText),
-    ('status', str, gtk.CellRendererText),
-    ('image', str, gtk.CellRendererText),
-    ('model', str, gtk.CellRendererText),
-    ('color', str, None),
-  ]
-  colorColIdx = len(listColumns)-1
-  _deviceColIdx = 0
-  _colorMap = {
-    'writing':          '#FFFF00',
-    'verifying':        '#6060FF',
-    'finished':         '#00FF00',
-    'error':            '#FF0000',
-  }
-
-  def __init__(self, model) :
-    self._model = model
-    tColoredListWindow.__init__(self, "%s: Write USB sticks" % (PROGNAME,))
-
-  def setStatus(self, device, status, image, model) :
-    row = self.findRow(device)
-    if row is None :
-      row = self.listModel.append()
-
-    color = self._colorMap[status]
-    self.listModel[row] = [locals()[name] for (name, t, r) in self.listColumns]
-
-  def findRow(self, device) :
-    return tColoredListWindow.findRow(self, self._deviceColIdx, device)
-
-
-class tModel (object) :
-
-  def __init__(self, imageprefix) :
-    self.imageprefix = imageprefix
-    self.mode = 'read'
-
-  def setMode(self, mode) :
-    if mode not in ['read', 'write'] :
-      raise ValueError("Unknown mode: %r" % (mode,))
-
-    self.mode = mode
-
-
-class tApp (object) :
-
-  def __init__(self) :
-    self._imageDir = self._chooseImageDir()
-    self._model = tModel(self._imageDir)
-    if self._askReadImage() :
-      self._win = tReadWindow(self._model)
-      self._win.writeButton.connect("clicked", self._showWriteWindow)
-    else :
-      self._win = None
-      self._showWriteWindow()
-
-  def _showWriteWindow(self, *args) :
-    if self._win :
-      # close read window
-      self._win.destroy()
-
-    self._win = tWriteWindow(self._model)
-    self._win.setStatus("/dev/sda", "writing", os.path.join(self._imageDir, "test-1gb.img"), "FOOBAR STICK")
-    self._win.setStatus("/dev/sdb", "verifying", os.path.join(self._imageDir, "test-1gb.img"), "SL STICK")
-    self._win.setStatus("/dev/sdc", "finished", os.path.join(self._imageDir, "test-1gb.img"), "SL STICK")
-    self._win.setStatus("/dev/sdd", "error", os.path.join(self._imageDir, "test-1gb.img"), "SL STICK")
-    self._win.connect('destroy', lambda *w: self._close())
-    self._model.setMode('write')
-
-  def _chooseImageDir(self) :
-    """
-    Query image directory from user. Will exit program in case user cancels.
-    """
-    imageChooserDialog = gtk.FileChooserDialog(title="Choose directory containing image files",
-      action=gtk.FILE_CHOOSER_ACTION_SELECT_FOLDER,
-      buttons=(gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_OPEN, gtk.RESPONSE_OK))
-    imageChooserDialog.set_default_response(gtk.RESPONSE_OK)
-    if imageChooserDialog.run() != gtk.RESPONSE_OK :
-      sys.exit(10)
-
-    dir = imageChooserDialog.get_filename()
-    imageChooserDialog.destroy()
-    return dir
-
-  def _askReadImage(self) :
-    """
-    Ask the user whether to create an image of a USB stick.
-    """
-    askDialog = gtk.MessageDialog(type=gtk.MESSAGE_QUESTION,
-      buttons=gtk.BUTTONS_NONE,
-      message_format="\n".join([x.strip() for x in """
-        Once the write mode is active, it will overwrite ANY USB stick you plug in. You now (and only now) have the chance to read an existing USB stick and save an image of it to the directory you just chose.
-        What do you want to do?
-        """.split("\n")]))
-
-    askDialog.add_buttons("Continue to write mode", gtk.RESPONSE_CANCEL, "Create image first", gtk.RESPONSE_OK)
-    res = askDialog.run()
-    askDialog.destroy()
-    return (res == gtk.RESPONSE_OK)
-
-  def run(self) :
-    gtk.main()
-
-  def _close(self) :
-    gtk.main_quit()
-
-
-def printSyntax(myName) :
-  print "Syntax: %s" % (myName,)
-  return 100
-
-def main(myName, args) :
-  if args :
-    return printSyntax(myName)
-
-  tApp().run()
-
-
-sys.exit(main(sys.argv[0], sys.argv[1:]))
diff --git a/soas-assimilator.py b/soas-assimilator.py
index ab90917..847a3b6 100755
--- a/soas-assimilator.py
+++ b/soas-assimilator.py
@@ -1,13 +1,18 @@
 #!/usr/bin/env python
 import os
-import shutil
 import sys
+import shutil
 import time
 import threading
 import traceback
-import dbus
-import dbus.mainloop.glib
 import gobject
+import gtk
+import dbus
+import dbus.glib # seems to automatically integrate dbus with GTK main loop
+
+
+PROGNAME = "soas-assimilator"
+
 
 try :
   import hashlib
@@ -18,17 +23,21 @@ except ImportError :
 
 
 _checksum_buffer_size = 1024*1024
-def checksum(f) :
+def checksum(f, max_size) :
   checksum = sha1()
-  buf = f.read(_checksum_buffer_size)
+  buf = f.read(min(_checksum_buffer_size, max_size))
   while buf :
     checksum.update(buf)
-    buf = f.read(_checksum_buffer_size)
+    max_size -= len(buf)
+    if not max_size :
+      break
+
+    buf = f.read(min(_checksum_buffer_size, max_size))
 
   return checksum.hexdigest()
 
-_sizeUnits = "KMGTP"
-def format_size(size) :
+_sizeUnits = " KMGTP"
+def formatSize(size) :
   if not size :
     return "0"
 
@@ -37,41 +46,113 @@ def format_size(size) :
     size //= 1024
     unitIdx += 1
 
-  return "%d%s" % (size, _sizeUnits[unitIdx])
+  return "%d%sB" % (size, _sizeUnits[unitIdx])
 
 
 
-class Logger (object) :
+class tColoredListWindow (gtk.Window) :
 
-  def __init__(self) :
-    self._lock = threading.Lock()
+  def __init__(self, title) :
+    gtk.Window.__init__(self)
+    self.set_title(title)
+    self.initWidgets()
 
-  def start(self, image_path, dev_path, image_size, model, dev_size) :
-    self._print("%s: starting to write %s (%s) on %s (%s)" % (dev_path, image_path, format_size(image_size), model, format_size(dev_size)))
+  def initWidgets(self) :
+    self.vbox = gtk.VBox(False, 10)
+    self.add(self.vbox)
+    self._scrolledList = gtk.ScrolledWindow()
+    self._scrolledList.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
+    self.vbox.pack_start(self._scrolledList)
+    self.listModel = gtk.ListStore(*[_type for (name, _type, r) in self.listColumns])
+    self._listView = gtk.TreeView(self.listModel)
+    self._scrolledList.add(self._listView)
+    self._addListColumns()
+    self.show_all()
 
-  def error(self, dev_path, exc) :
-    self._print("%s: aborted with error: %s" % (dev_path, exc))
+  def _addListColumns(self) :
+    for (idx, (name, _type, renderer)) in enumerate(self.listColumns) :
+      if not renderer :
+        continue
 
-  def verify(self, dev_path, time_delta, image_size) :
-    self._print("%s: starting to verify (write took %ds, %.3fMB/s)" % (dev_path,
-      time_delta,float(image_size)/1024/1024/time_delta))
+      r = renderer()
+      r.set_property('background-set', True)
+      if (_type == bool) :
+        col = gtk.TreeViewColumn(name, r, active=idx, background=self.colorColIdx)
+      else :
+        col = gtk.TreeViewColumn(name, r, text=idx, background=self.colorColIdx)
+        col.set_resizable(True)
+
+      col.set_sort_column_id(idx)
+      self._listView.append_column(col)
+
+  def findRow(self, col, content) :
+    for idx in range(len(self.listModel)) :
+      if (self.listModel[idx][col] == content) :
+        return idx
+
+
+class tStatusListWindow (tColoredListWindow) :
+
+  listColumns = [
+    ('device', str, gtk.CellRendererText),
+    ('status', str, gtk.CellRendererText),
+    ('image',  str, gtk.CellRendererText),
+    ('model',  str, gtk.CellRendererText),
+    ('info',   str, gtk.CellRendererText),
+    ('color',  str, None),
+  ]
+  _deviceColIdx = 0
+  colorColIdx = len(listColumns)-1
+  _colorMap = {
+    'reading':   '#FFFF00',
+    'writing':   '#FFFF00',
+    'verifying': '#6060FF',
+    'success':   '#00FF00',
+    'error':     '#FF0000',
+  }
+
+  def __init__(self, title) :
+    tColoredListWindow.__init__(self, "%s: %s" % (PROGNAME, title))
+
+  def setStatus(self, device, info) :
+    row = self.findRow(device)
+    if row is None :
+      row = self.listModel.append()
+
+    info = dict(info)
+    info['color'] = self._colorMap[info['status']]
+    self.listModel[row] = [info.get(name, '') for (name, t, r) in self.listColumns]
+
+  def findRow(self, device) :
+    return tColoredListWindow.findRow(self, self._deviceColIdx, device)
+
+
+class tReadWindow (tStatusListWindow) :
 
-  def success(self, dev_path, time_delta, image_size) :
-    self._print("%s: finished successfully (verify took %ds, %.3fMB/s)" % (dev_path,
-      time_delta,float(image_size)/1024/1024/time_delta))
+  def __init__(self) :
+    tColoredListWindow.__init__(self, "Read USB sticks")
+
+  def initWidgets(self) :
+    tStatusListWindow.initWidgets(self)
+    self._hButtonBox = gtk.HButtonBox()
+    self.writeButton = gtk.Button("Continue to write mode")
+    self.writeButton.modify_bg(gtk.STATE_NORMAL, gtk.gdk.color_parse("#F00000"))
+    self.writeButton.modify_bg(gtk.STATE_PRELIGHT, gtk.gdk.color_parse("#FF0000"))
+    self.writeButton.show()
+    self._hButtonBox.pack_start(self.writeButton)
+    self._hButtonBox.show()
+    self.vbox.pack_end(self._hButtonBox, expand=False)
 
-  def _print(self, msg) :
-    self._lock.acquire()
-    try :
-      print msg
 
-    finally :
-      self._lock.release()
+class tWriteWindow (tStatusListWindow) :
+
+  def __init__(self) :
+    tStatusListWindow.__init__(self, "Write USB sticks")
 
 
 class CopyThread (threading.Thread) :
 
-  def __init__(self, image_path, dev_path, model, dev_size, image_size, checksum, logger) :
+  def __init__(self, image_path, dev_path, model, dev_size, image_size, checksum, status_cb) :
     threading.Thread.__init__(self)
     self.image_path = image_path
     self.dev_path = dev_path
@@ -79,44 +160,46 @@ class CopyThread (threading.Thread) :
     self.image_size = image_size
     self.dev_size = dev_size
     self.checksum = checksum
-    self.logger = logger
+    self.status_cb = status_cb
 
   def run(self) :
     self.start_time = time.time()
-    self.logger.start(self.image_path, self.dev_path, self.image_size, self.model, self.dev_size)
+    self.status_cb(device=self.dev_path, status='writing', image=self.image_path, image_size=self.image_size, model=self.model, device_size=self.dev_size)
     try :
       image_f, dev_f = file(self.image_path, "rb"), file(self.dev_path, "wb+")
 
       try :
         # copy
         shutil.copyfileobj(image_f, dev_f)
+        dev_f.flush()
         dev_f.seek(0)
   
         cur_time = time.time()
-        self.logger.verify(self.dev_path, cur_time - self.start_time, self.image_size)
+        self.status_cb(device=self.dev_path, status="verifying", write_time=(cur_time - self.start_time))
         self.start_time = cur_time
   
         # verify
-        dev_checksum = checksum(dev_f)
+        dev_checksum = checksum(dev_f, self.image_size)
         if (dev_checksum != self.checksum) :
-          return self.logger.error(self.dev_path,
-            "verification failed: %r (image) != %r (device)" % (self.checksum, dev_checksum))
+          return self.status_cb(device=self.dev_path, status='error',
+            error="verification failed: %r (image) != %r (device)" % (self.checksum, dev_checksum))
 
       finally :
         image_f.close(), dev_f.close()
 
     except :
-      self.logger.error(self.dev_path, sys.exc_info()[1])
+      self.status_cb(device=self.dev_path, status='error', error=sys.exc_info()[1])
 
     else :
-      self.logger.success(self.dev_path, time.time() - self.start_time, self.image_size)
+      self.status_cb(device=self.dev_path, status='success', verify_time=(time.time() - self.start_time))
 
 
-class DeviceListener (object) :
+class tDeviceListener (object) :
 
-  def __init__(self, image_dir) :
+  def __init__(self, image_dir, status_cb) :
+    self.mode = 'read'
     self.image_dir = image_dir
-    self.logger = Logger()
+    self.status_cb = status_cb
     self._images_dict = {}
     self._images_flat = []
     self._images_lastupdate = 0
@@ -127,8 +210,11 @@ class DeviceListener (object) :
                                  "org.freedesktop.Hal",
                                  "/org/freedesktop/Hal/Manager")
 
-  def run(self) :
-    gobject.MainLoop().run()
+  def setMode(self, mode) :
+    if mode not in ['read', 'write'] :
+      raise ValueError("Unknown mode: %r" % (mode,))
+
+    self.mode = mode
 
   def hal_device_added(self, udi) :
     try :
@@ -145,8 +231,9 @@ class DeviceListener (object) :
 #      for (k,v) in props.items() :
 #        print "%30s %s" % (k,v)
 
-      self.assimilate(str(props['storage.model']), str(props['block.device']),
-        int(props['storage.removable.media_size']))
+      if self.mode == 'write' :
+        self.assimilate(str(props['storage.model']), str(props['block.device']),
+          int(props['storage.removable.media_size']))
 
     except :
       # don't break if we have trouble with a single device, but show the error
@@ -157,9 +244,9 @@ class DeviceListener (object) :
       image_size, (image_path, image_checksum) = self._find_image(dev_size)
 
     except ValueError :
-      return self.logger.error(dev_path, "No matching image found")
+      return self.status_cb(dev_path, model=model, device_size=dev_size, status="error", error="No matching image found")
 
-    thread = CopyThread(image_path, dev_path, model, dev_size, image_size, image_checksum, self.logger)
+    thread = CopyThread(image_path, dev_path, model, dev_size, image_size, image_checksum, self.status_cb)
     thread.start()
 
   def _find_image(self, size) :
@@ -195,22 +282,97 @@ class DeviceListener (object) :
     self._images_lastupdate = curTime
 
 
+class tApp (object) :
+
+  def __init__(self) :
+    self._imageDir = self._chooseImageDir()
+    self._devStatus = {}
+    self._deviceListener = tDeviceListener(self._imageDir, self._updateStatus)
+    if self._askReadImage() :
+      self._win = tReadWindow()
+      self._win.writeButton.connect("clicked", self._showWriteWindow)
+    else :
+      self._win = None
+      self._showWriteWindow()
+
+  def _showWriteWindow(self, *args) :
+    if self._win :
+      # close read window
+      self._win.destroy()
+
+    self._win = tWriteWindow()
+    self._win.connect('destroy', lambda *w: self._close())
+    self._deviceListener.setMode('write')
+
+  def _updateStatus(self, device, status, **kwArgs) :
+    devStatus = self._devStatus.setdefault(device, {'device': device})
+    devStatus.update(kwArgs, status=status)
+    info = {'device': device, 'status': status,
+        'model': "%s (%s)" % (devStatus['model'], formatSize(devStatus['device_size'])),
+        'image': "%s (%s)" % (devStatus.get('image', ''), formatSize(devStatus.get('image_size', 0))),
+        'info': ''}
+
+    if (status in ['verifying', 'success']) :
+      info['info'] = "writing took %ds, %.3fMB/s" % (devStatus['write_time'],
+        float(devStatus['image_size'])/1024/1024/devStatus['write_time'])
+
+    if (status == 'success') :
+      info['info'] += ", verification took %ds, %.3fMB/s" % (devStatus['verify_time'],
+        float(devStatus['image_size'])/1024/1024/devStatus['verify_time'])
+
+    elif (status == 'error') :
+      info['info'] = devStatus['error']
+
+    self._win.setStatus(device, info)
+
+  def _chooseImageDir(self) :
+    """
+    Query image directory from user. Will exit program in case user cancels.
+    """
+    imageChooserDialog = gtk.FileChooserDialog(title="Choose directory containing image files",
+      action=gtk.FILE_CHOOSER_ACTION_SELECT_FOLDER,
+      buttons=(gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_OPEN, gtk.RESPONSE_OK))
+    imageChooserDialog.set_default_response(gtk.RESPONSE_OK)
+    if imageChooserDialog.run() != gtk.RESPONSE_OK :
+      sys.exit(10)
+
+    dir = imageChooserDialog.get_filename()
+    imageChooserDialog.destroy()
+    return dir
+
+  def _askReadImage(self) :
+    """
+    Ask the user whether to create an image of a USB stick.
+    """
+    askDialog = gtk.MessageDialog(type=gtk.MESSAGE_QUESTION,
+      buttons=gtk.BUTTONS_NONE,
+      message_format="\n".join([x.strip() for x in """
+        Once the write mode is active, it will overwrite ANY USB stick you plug in. You now (and only now) have the chance to read an existing USB stick and save an image of it to the directory you just chose.
+        What do you want to do?
+        """.split("\n")]))
+
+    askDialog.add_buttons("Continue to write mode", gtk.RESPONSE_CANCEL, "Create image first", gtk.RESPONSE_OK)
+    res = askDialog.run()
+    askDialog.destroy()
+    return (res == gtk.RESPONSE_OK)
+
+  def run(self) :
+    gtk.main()
+
+  def _close(self) :
+    gtk.main_quit()
+
 
 def printSyntax(myName) :
-  print "Syntax: %s <image dir>" % (myName,)
-  print "Sample: %s ~/Soas-20090324" % (myName,)
+  print "Syntax: %s" % (myName,)
   return 100
 
 def main(myName, args) :
-  if (len(args) < 1) :
+  if args :
     return printSyntax(myName)
 
-  image_dir, = args
-
-  dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
   gobject.threads_init()
-  DeviceListener(image_dir).run()
+  tApp().run()
 
-if __name__ == '__main__' :
-  sys.exit(main(sys.argv[0], sys.argv[1:]))
 
+sys.exit(main(sys.argv[0], sys.argv[1:]))