Add state-safe-copy.py

state-safe-copy.py

+#!/usr/bin/env python2
+# -*- coding: utf-8 -*-
+"""
+Usage: python cros-safe-copy.py SRC DST [SLEEP]
+
+Where SRC is the source file path
+      DST is the destination file path
+      SLEEP is the time to wait between two copied chunk (in second)
+
+Do a copy of flushed everytime chunks from SRC to DST. This is usefull if your
+`dd` command makes your machine rebooting. More, this script save the state
+during copying so that even if the machine reboots, you can copy from where it
+was stopped. (Be careful to run your script where CWD is persistent! Avoid /tmp
+or ~/Downloads on Chromebooks!)
+
+Example:
+  sudo python cros-safe-copy.py NayuOS.img /dev/sda
+  # savestate.json file is saved in CWD
+"""
+
+import sys, os, time, json
+
+input = sys.argv[1]
+output = sys.argv[2]
+sleep = int(sys.argv[3]) if len(sys.argv) > 3 else 0
+
+def saveState(d):
+  open("savestate.json.lock", "wx").write(json.dumps(d))
+  os.rename("savestate.json.lock", "savestate.json")
+
+def loadState():
+  try:
+    state = json.loads(open("savestate.json").read())
+    print "Using savestate.json"
+    return state
+  except IOError:
+    print "Creating savestate.json"
+    MiB = 1024*1024
+    bs = 65536
+    saveState(dict(
+      input=input,
+      output=output,
+      bs=bs,
+      count=50*MiB/bs,
+      skip=0,
+      seek=0,
+    ))
+    return loadState()
+
+def quote(args):
+  return " ".join(["'%s'" % a.replace("'", "'\\''") for a in args])
+
+state = loadState()
+assert state["input"] == input
+assert state["output"] == output
+
+while state["skip"] * state["bs"] < os.stat(input).st_size:
+  command = quote([
+    "dd",
+    "bs=%d" % state["bs"],
+    "count=%d" % state["count"],
+    "seek=%d" % state["seek"],
+    "skip=%d" % state["skip"],
+    "if=%s" % input,
+    "of=%s" % output,
+  ])
+  print "$", command
+  if os.system(command):
+    break
+  print "$ sync"
+  if os.system("sync"):  # seems not useful
+    break
+  state["skip"] += state["count"]
+  state["seek"] += state["count"]
+  saveState(state)
+  if sleep > 0:
+    print "sleeping for %d second(s)" % sleep
+    time.sleep(sleep)