blocksync.py

#!/usr/bin/env python
"""
Synchronise block devices over the network

Copyright 2006-2008 Justin Azoff <[email protected]>
Copyright 2011 Robert Coup <[email protected]>
Copyright 2012 Holger Ernst <[email protected]>
Copyright 2014 Robert McQueen <[email protected]>
License: GPL

Getting started:

* Copy blocksync.py to the home directory on the remote host & make it executable
* Make sure your remote user is either root or can sudo (use -s for sudo)
* Make sure your local user can ssh to the remote host (use -i for a SSH key)
* Invoke:
    python blocksync.py /dev/source [user@]remotehost [/dev/dest]

* Specify localhost for local usage:
    python blocksync.py /dev/source localhost /dev/dest
"""

import os
import sys
from zlib import adler32
import subprocess
import time

SAME = "same\n"
DIFF = "diff\n"


def do_open(f, mode):
    f = open(f, mode)
    f.seek(0, 2)
    size = f.tell()
    f.seek(0)
    return f, size


def getblocks(f, blocksize):
    while 1:
        block = f.read(blocksize)
        if not block:
            break
        yield block


def server(dev, blocksize):
    print dev, blocksize
    f, size = do_open(dev, 'r+')
    print size
    sys.stdout.flush()

    for block in getblocks(f, blocksize):
        print "%08x" % (adler32(block) & 0xFFFFFFFF)
        sys.stdout.flush()
        res = sys.stdin.readline()
        if res != SAME:
            newblock = sys.stdin.read(blocksize)
            f.seek(-len(newblock), 1)
            f.write(newblock)


def sync(srcdev, dsthost, dstdev=None, blocksize=1024 * 1024, keyfile=None, pause=0, sudo=False):

    if not dstdev:
        dstdev = srcdev

    print "Block size is %0.1f MB" % (float(blocksize) / (1024 * 1024))

    pause_ms = 0
    if pause:
        # sleep() wants seconds...
        pause_ms = float(pause) / 1000
        print "Slowing down for %d ms/block (%0.4f sec/block)" % (pause, pause_ms)

    cmd = []
    if dsthost != 'localhost':
        cmd += ['ssh', '-c', 'blowfish']
        if keyfile:
            cmd += ['-i', keyfile]
        cmd += [dsthost]
    if sudo:
        cmd += ['sudo']
    cmd += ['python', 'blocksync.py', 'server', dstdev, '-b', str(blocksize)]

    print "Running: %s" % " ".join(cmd)

    p = subprocess.Popen(cmd, bufsize=0, stdin=subprocess.PIPE, stdout=subprocess.PIPE, close_fds=True)
    p_in, p_out = p.stdin, p.stdout

    line = p_out.readline()
    p.poll()
    if p.returncode is not None:
        print "Error connecting to or invoking blocksync on the remote host!"
        sys.exit(1)

    a, b = line.split()
    if a != dstdev:
        print "Dest device (%s) doesn't match with the remote host (%s)!" % (dstdev, a)
        sys.exit(1)
    if int(b) != blocksize:
        print "Source block size (%d) doesn't match with the remote host (%d)!" % (blocksize, int(b))
        sys.exit(1)

    try:
        f, size = do_open(srcdev, 'r')
    except Exception, e:
        print "Error accessing source device! %s" % e
        sys.exit(1)

    line = p_out.readline()
    p.poll()
    if p.returncode is not None:
        print "Error accessing device on remote host!"
        sys.exit(1)
    remote_size = int(line)
    if size > remote_size:
        print "Source device size (%d) doesn't fit into remote device size (%d)!" % (size, remote_size)
        sys.exit(1)
    elif size < remote_size:
        print "Source device size (%d) is smaller than remote device size (%d), proceeding anyway" % (size, remote_size)

    same_blocks = diff_blocks = 0
    interactive = os.isatty(sys.stdout.fileno())

    print "Starting sync..."
    t0 = time.time()
    t_last = t0
    size_blocks = size / blocksize
    for i, l_block in enumerate(getblocks(f, blocksize)):
        l_sum = "%08x" % (adler32(l_block) & 0xFFFFFFFF)
        r_sum = p_out.readline().strip()
        if pause_ms:
            time.sleep(pause_ms)

        if l_sum == r_sum:
            p_in.write(SAME)
            p_in.flush()
            same_blocks += 1
        else:
            p_in.write(DIFF)
            p_in.flush()
            p_in.write(l_block)
            p_in.flush()
            diff_blocks += 1

        t1 = time.time()
        rate = (i + 1.0) * blocksize / (1024.0 * 1024.0) / (t1 - t0)

        if not interactive:
            continue

        if t1 - t_last > 1 or (same_blocks + diff_blocks) >= size_blocks:
            print "\rsame: %d, diff: %d, %d/%d, %5.1f MB/s" % (same_blocks, diff_blocks, same_blocks + diff_blocks, size_blocks, rate),
            t_last = t1

    if interactive:
        print
    else:
        print "same: %d, diff: %d, %d/%d, %5.1f MB/s" % (same_blocks, diff_blocks, same_blocks + diff_blocks, size_blocks, rate)

    print "\nCompleted in %d seconds" % (time.time() - t0)

    return same_blocks, diff_blocks

if __name__ == "__main__":
    from optparse import OptionParser
    parser = OptionParser(usage="%prog [options] /dev/source [user@]remotehost [/dev/dest]")
    parser.add_option("-b", "--blocksize", dest="blocksize", type="int", help="block size (bytes, defaults to 1MB)", default=1024 * 1024)
    parser.add_option("-i", "--id", dest="keyfile", help="ssh public key file")
    parser.add_option("-p", "--pause", dest="pause", type="int", help="pause between processing blocks, reduces system load (ms, defaults to 0)")
    parser.add_option("-s", "--sudo", dest="sudo", action="store_true", help="use sudo on the remote end (defaults to off)", default=False)
    (options, args) = parser.parse_args()

    if len(args) < 2:
        parser.print_help()
        print __doc__
        sys.exit(1)

    if args[0] == 'server':
        dstdev = args[1]
        server(dstdev, options.blocksize)
    else:
        srcdev = args[0]
        dsthost = args[1]
        if len(args) > 2:
            dstdev = args[2]
        else:
            dstdev = None
        sync(srcdev, dsthost, dstdev, options.blocksize, options.keyfile, options.pause, options.sudo)