first working prototype on a LiveCD

1 files changed, 160 insertions, 0 deletions

diff --git a/soas-assimilator.py b/soas-assimilator.py
new file mode 100755
index 0000000..37106f9
--- /dev/null
+++ b/soas-assimilator.py
@@ -0,0 +1,160 @@
+#!/usr/bin/env python
+import os
+import shutil
+import sys
+import time
+import threading
+import traceback
+import dbus
+import dbus.mainloop.glib
+import gobject
+
+
+class Logger (object) :
+
+  def __init__(self) :
+    self._lock = threading.Lock()
+
+  def start(self, image_path, dev_path, model) :
+    self._print("%s: starting to write %s on %s" % (dev_path, image_path, model))
+
+  def error(self, dev_path, exc) :
+    self._print("%s: aborted with error: %s" % (dev_path, exc))
+
+  def success(self, dev_path, time_delta, size) :
+    self._print("%s: finished successfully (%ds, %.3fMB/s)" % (dev_path,
+      time_delta,float(size)/1024/1024/time_delta))
+
+  def _print(self, msg) :
+    self._lock.acquire()
+    try :
+      print msg
+
+    finally :
+      self._lock.release()
+
+
+class CopyThread (threading.Thread) :
+
+  def __init__(self, image_path, dev_path, model, size, logger) :
+    threading.Thread.__init__(self)
+    self.image_path = image_path
+    self.dev_path = dev_path
+    self.model = model
+    self.size = size
+    self.logger = logger
+
+  def run(self) :
+    self.start_time = time.time()
+    self.logger.start(self.image_path, self.dev_path, self.model)
+    try :
+      shutil.copyfileobj(file(self.image_path, "rb"), file(self.dev_path, "wb"))
+
+    except :
+      self.logger.error(self.dev_path, sys.exc_info()[1])
+
+    else :
+      self.logger.success(self.dev_path, time.time() - self.start_time, self.size)
+
+
+class DeviceListener (object) :
+
+  def __init__(self, image_prefix) :
+    self.image_prefix = image_prefix
+    self.logger = Logger()
+    self._images_dict = {}
+    self._images_flat = []
+    self._images_lastupdate = 0
+    self.bus = dbus.Bus(dbus.Bus.TYPE_SYSTEM)
+    self.bus.add_signal_receiver(self.hal_device_added,
+                                 "DeviceAdded",
+                                 "org.freedesktop.Hal.Manager",
+                                 "org.freedesktop.Hal",
+                                 "/org/freedesktop/Hal/Manager")
+
+  def run(self) :
+    gobject.MainLoop().run()
+
+  def hal_device_added(self, udi) :
+    try :
+      proxy = self.bus.get_object("org.freedesktop.Hal", udi)
+      dev = dbus.Interface(proxy, "org.freedesktop.Hal.Device")
+      props = dev.GetAllProperties()
+      if   (props.get("info.category") != "storage") \
+        or (props.get("storage.removable") != 1) \
+        or (props.get("storage.drive_type") != "disk") :
+        # not a USB stick
+        return
+
+      # for debugging
+#      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']))
+
+    except :
+      # don't break if we have trouble with a single device, but show the error
+      traceback.print_exc()
+
+  def assimilate(self, model, dev_path, size) :
+#    image_path = "%s-%s.img" % (self.image_prefix, self._format_size(size))
+    image_path = self._find_image(size)
+    thread = CopyThread(image_path, dev_path, model, size, self.logger)
+    thread.start()
+
+  _sizeUnits = "KMGTP"
+  def _format_size(self, size) :
+    if not size :
+      return "0"
+
+    unitIdx = 0
+    while ((size % 1024) == 0) :
+      size //= 1024
+      unitIdx += 1
+
+    return "%d%s" % (size, self._sizeUnits[unitIdx])
+
+  def _find_image(self, size) :
+    """
+    Try to find image to put on current stick. Will return the largest one
+    still fitting on the stick.
+    """
+    self._cache_images()
+    return [iname for (isize, iname) in self._images_flat if (isize <= size)][-1]
+
+
+  _images_cachetime = 10 # rescan images every 10 seconds
+  def _cache_images(self) :
+    curTime = time.time()
+    if (curTime - self._images_lastupdate) < self._images_cachetime :
+      return
+
+    image_names = [fname for fname in os.listdir(".") if fname.startswith(self.image_prefix) and fname.endswith(".img")]
+    for name in image_names :
+      if name not in self._images_dict :
+        self._images_dict[os.stat(name).st_size] = name
+        self._images_flat = sorted(self._images_dict.items())
+
+    self._images_lastupdate = curTime
+
+
+
+def printSyntax(myName) :
+  print "Syntax: %s <image prefix>" % (myName,)
+  print "Sample: %s Soas-20090324" % (myName,)
+  return 100
+
+def main(myName, args) :
+  if (len(args) < 1) :
+    return printSyntax(myName)
+
+  image_prefix, = args
+
+  dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
+  gobject.threads_init()
+  DeviceListener(image_prefix).run()
+
+if __name__ == '__main__' :
+  sys.exit(main(sys.argv[0], sys.argv[1:]))
+