finish read mode implementation

1 files changed, 78 insertions, 18 deletions

diff --git a/soas-assimilator.py b/soas-assimilator.py
index 847a3b6..d862ad4 100755
--- a/soas-assimilator.py
+++ b/soas-assimilator.py
@@ -150,7 +150,54 @@ class tWriteWindow (tStatusListWindow) :
     tStatusListWindow.__init__(self, "Write USB sticks")
 
 
-class CopyThread (threading.Thread) :
+class ReadThread (threading.Thread) :
+
+  _bufferSize = 1024*1024
+
+  def __init__(self, dev_path, model, dev_size, image_prefix, status_cb) :
+    threading.Thread.__init__(self)
+    self.dev_path = dev_path
+    self.model = model
+    self.dev_size = dev_size
+    self.image_prefix = image_prefix
+    self.image_path = image_prefix + ".img"
+    self.checksum_path = image_prefix + ".sha1"
+    self.status_cb = status_cb
+
+  def run(self) :
+    start_time = time.time()
+    self.status_cb(device=self.dev_path, status='reading', image=self.image_path, image_size=self.dev_size, model=self.model, device_size=self.dev_size)
+    try :
+      imageF, devF = file(self.image_path, "wb"), file(self.dev_path, "rb")
+
+      try :
+        # copy
+        checksum = self._copy(imageF, devF)
+        imageF.flush()
+
+      finally :
+        imageF.close(), devF.close()
+
+      file(self.checksum_path, "w").write("%s\n" % (checksum,))
+
+    except :
+      self.status_cb(device=self.dev_path, status='error', error=sys.exc_info()[1])
+
+    else :
+      self.status_cb(device=self.dev_path, status='success', read_time=(time.time() - start_time))
+
+  def _copy(self, imageF, devF) :
+    checksum = sha1()
+    buf = devF.read(self._bufferSize)
+    while buf :
+      checksum.update(buf)
+      imageF.write(buf)
+      buf = devF.read(self._bufferSize)
+
+    return checksum.hexdigest()
+
+
+class WriteThread (threading.Thread) :
 
   def __init__(self, image_path, dev_path, model, dev_size, image_size, checksum, status_cb) :
     threading.Thread.__init__(self)
@@ -163,35 +210,35 @@ class CopyThread (threading.Thread) :
     self.status_cb = status_cb
 
   def run(self) :
-    self.start_time = time.time()
+    start_time = time.time()
     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+")
+      imageF, devF = 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)
+        shutil.copyfileobj(imageF, devF)
+        devF.flush()
+        devF.seek(0)
   
         cur_time = time.time()
-        self.status_cb(device=self.dev_path, status="verifying", write_time=(cur_time - self.start_time))
-        self.start_time = cur_time
+        self.status_cb(device=self.dev_path, status="verifying", write_time=(cur_time - start_time))
+        start_time = cur_time
   
         # verify
-        dev_checksum = checksum(dev_f, self.image_size)
+        dev_checksum = checksum(devF, self.image_size)
         if (dev_checksum != self.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()
+        imageF.close(), devF.close()
 
     except :
       self.status_cb(device=self.dev_path, status='error', error=sys.exc_info()[1])
 
     else :
-      self.status_cb(device=self.dev_path, status='success', verify_time=(time.time() - self.start_time))
+      self.status_cb(device=self.dev_path, status='success', verify_time=(time.time() - start_time))
 
 
 class tDeviceListener (object) :
@@ -234,6 +281,9 @@ class tDeviceListener (object) :
       if self.mode == 'write' :
         self.assimilate(str(props['storage.model']), str(props['block.device']),
           int(props['storage.removable.media_size']))
+      elif self.mode == 'read' :
+        self.readImage(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
@@ -246,7 +296,12 @@ class tDeviceListener (object) :
     except ValueError :
       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.status_cb)
+    thread = WriteThread(image_path, dev_path, model, dev_size, image_size, image_checksum, self.status_cb)
+    thread.start()
+
+  def readImage(self, model, dev_path, dev_size) :
+    image_prefix = "dump-%s" % (formatSize(dev_size),)
+    thread = ReadThread(dev_path, model, dev_size, os.path.join(self.image_dir, image_prefix), self.status_cb)
     thread.start()
 
   def _find_image(self, size) :
@@ -262,7 +317,6 @@ class tDeviceListener (object) :
       # no matching image found
       raise ValueError("No matching image found")
 
-
   _images_cachetime = 10 # rescan images every 10 seconds
   def _cache_images(self) :
     curTime = time.time()
@@ -291,17 +345,19 @@ class tApp (object) :
     if self._askReadImage() :
       self._win = tReadWindow()
       self._win.writeButton.connect("clicked", self._showWriteWindow)
+      self._destroyCbHandle = self._win.connect('destroy', lambda *w: self._close())
     else :
       self._win = None
       self._showWriteWindow()
 
   def _showWriteWindow(self, *args) :
     if self._win :
-      # close read window
+      # close read window, but prevent closing the application
+      self._win.handler_disconnect(self._destroyCbHandle)
       self._win.destroy()
 
     self._win = tWriteWindow()
-    self._win.connect('destroy', lambda *w: self._close())
+    self._destroyCbHandle = self._win.connect('destroy', lambda *w: self._close())
     self._deviceListener.setMode('write')
 
   def _updateStatus(self, device, status, **kwArgs) :
@@ -313,10 +369,14 @@ class tApp (object) :
         '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 (self._deviceListener.mode == "read") :
+        info['info'] = "reading took %ds, %.3fMB/s" % (devStatus['read_time'],
+          float(devStatus['image_size'])/1024/1024/devStatus['read_time'])
+      else :
+        info['info'] = "writing took %ds, %.3fMB/s" % (devStatus['write_time'],
+          float(devStatus['image_size'])/1024/1024/devStatus['write_time'])
 
-    if (status == 'success') :
+    if (status == 'success') and (self._deviceListener.mode == "write") :
       info['info'] += ", verification took %ds, %.3fMB/s" % (devStatus['verify_time'],
         float(devStatus['image_size'])/1024/1024/devStatus['verify_time'])