soas-assimilator.py: implement verify ; use image directory instead of file prefix

1 files changed, 92 insertions, 36 deletions

diff --git a/soas-assimilator.py b/soas-assimilator.py
index 37106f9..ab90917 100755
--- a/soas-assimilator.py
+++ b/soas-assimilator.py
@@ -9,21 +9,56 @@ import dbus
 import dbus.mainloop.glib
 import gobject
 
+try :
+  import hashlib
+  sha1 = hashlib.sha1
+except ImportError :
+  import sha
+  sha1 = sha.new
+
+
+_checksum_buffer_size = 1024*1024
+def checksum(f) :
+  checksum = sha1()
+  buf = f.read(_checksum_buffer_size)
+  while buf :
+    checksum.update(buf)
+    buf = f.read(_checksum_buffer_size)
+
+  return checksum.hexdigest()
+
+_sizeUnits = "KMGTP"
+def format_size(size) :
+  if not size :
+    return "0"
+
+  unitIdx = 0
+  while ((size % 1024) == 0) :
+    size //= 1024
+    unitIdx += 1
+
+  return "%d%s" % (size, _sizeUnits[unitIdx])
+
+
 
 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 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 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 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))
+
+  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 _print(self, msg) :
     self._lock.acquire()
@@ -36,31 +71,51 @@ class Logger (object) :
 
 class CopyThread (threading.Thread) :
 
-  def __init__(self, image_path, dev_path, model, size, logger) :
+  def __init__(self, image_path, dev_path, model, dev_size, image_size, checksum, logger) :
     threading.Thread.__init__(self)
     self.image_path = image_path
     self.dev_path = dev_path
     self.model = model
-    self.size = size
+    self.image_size = image_size
+    self.dev_size = dev_size
+    self.checksum = checksum
     self.logger = logger
 
   def run(self) :
     self.start_time = time.time()
-    self.logger.start(self.image_path, self.dev_path, self.model)
+    self.logger.start(self.image_path, self.dev_path, self.image_size, self.model, self.dev_size)
     try :
-      shutil.copyfileobj(file(self.image_path, "rb"), file(self.dev_path, "wb"))
+      image_f, dev_f = file(self.image_path, "rb"), file(self.dev_path, "wb+")
+
+      try :
+        # copy
+        shutil.copyfileobj(image_f, dev_f)
+        dev_f.seek(0)
+  
+        cur_time = time.time()
+        self.logger.verify(self.dev_path, cur_time - self.start_time, self.image_size)
+        self.start_time = cur_time
+  
+        # verify
+        dev_checksum = checksum(dev_f)
+        if (dev_checksum != self.checksum) :
+          return self.logger.error(self.dev_path,
+            "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])
 
     else :
-      self.logger.success(self.dev_path, time.time() - self.start_time, self.size)
+      self.logger.success(self.dev_path, time.time() - self.start_time, self.image_size)
 
 
 class DeviceListener (object) :
 
-  def __init__(self, image_prefix) :
-    self.image_prefix = image_prefix
+  def __init__(self, image_dir) :
+    self.image_dir = image_dir
     self.logger = Logger()
     self._images_dict = {}
     self._images_flat = []
@@ -97,23 +152,15 @@ class DeviceListener (object) :
       # 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"
+  def assimilate(self, model, dev_path, dev_size) :
+    try :
+      image_size, (image_path, image_checksum) = self._find_image(dev_size)
 
-    unitIdx = 0
-    while ((size % 1024) == 0) :
-      size //= 1024
-      unitIdx += 1
+    except ValueError :
+      return self.logger.error(dev_path, "No matching image found")
 
-    return "%d%s" % (size, self._sizeUnits[unitIdx])
+    thread = CopyThread(image_path, dev_path, model, dev_size, image_size, image_checksum, self.logger)
+    thread.start()
 
   def _find_image(self, size) :
     """
@@ -121,7 +168,12 @@ class DeviceListener (object) :
     still fitting on the stick.
     """
     self._cache_images()
-    return [iname for (isize, iname) in self._images_flat if (isize <= size)][-1]
+    try :
+      return [(isize, image) for (isize, image) in self._images_flat if (isize <= size)][-1]
+
+    except IndexError :
+      # no matching image found
+      raise ValueError("No matching image found")
 
 
   _images_cachetime = 10 # rescan images every 10 seconds
@@ -130,30 +182,34 @@ class DeviceListener (object) :
     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")]
+    image_names = [os.path.join(self.image_dir, fname)
+      for fname in os.listdir(self.image_dir) if 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())
+      checksum_fname = name[:-3]+"sha1"
+      if (name in self._images_dict) or (not os.path.exists(checksum_fname)) :
+        continue
+
+      self._images_dict[os.stat(name).st_size] = (name,file(checksum_fname).read().strip())
+      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,)
+  print "Syntax: %s <image dir>" % (myName,)
+  print "Sample: %s ~/Soas-20090324" % (myName,)
   return 100
 
 def main(myName, args) :
   if (len(args) < 1) :
     return printSyntax(myName)
 
-  image_prefix, = args
+  image_dir, = args
 
   dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)
   gobject.threads_init()
-  DeviceListener(image_prefix).run()
+  DeviceListener(image_dir).run()
 
 if __name__ == '__main__' :
   sys.exit(main(sys.argv[0], sys.argv[1:]))