#!/usr/bin/env python
#
"""
== blocksync v1.4 (2023-11-09) == sync block devices over network ==

Getting started:

 * blocksync must be installed on the target host
 * make sure your remote user can either sudo or is root itself
 * make sure your local user can ssh to the remote host (not needed
   if remote is localhost)
 * run the blocksync command from source system
"""

import sys
from zlib import adler32
import subprocess
import time

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

def_blocksize = 1024 * 1024
def_cipher = "aes192-cbc"


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, f_exec):
    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)
	    if f_exec:
                f.seek(-len(newblock), 1)
                f.write(newblock)


def sync(srcdev, dsthost, dstdev=None, blocksize=def_blocksize, cipher=def_cipher, f_exec=False, f_sudo=True):

    if not dstdev:
        dstdev = srcdev

    exec_flag = "-n"
    if f_exec:
        exec_flag = "-x"

    print ""
    print " block size is %0.1f MiB" % (float(blocksize) / (1024 * 1024))

    if dsthost == "localhost":
        cmd = []
        f_sudo = False
    else:
        cmd = [ 'ssh', '-c', cipher, dsthost ]

    if f_sudo:
    	cmd.append( 'sudo' )

    cmd.extend( [ 'blocksync', 'server', dstdev, '-b', str(blocksize), exec_flag ] )
    print " starting server: %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 match remote device size (%d)!" % (size, remote_size)
        sys.exit(1)

    same_blocks = diff_blocks = 0

    print " starting sync ..."
    print ""

    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 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()
        if t1 - t_last > 1 or (same_blocks + diff_blocks) >= size_blocks:
            rate = (i + 1.0) * blocksize / (1024.0 * 1024.0) / (t1 - t0)
            print "\r  same: %d, diff: %d, %d/%d, %5.1f MiB/s" % (same_blocks, diff_blocks, same_blocks + diff_blocks, size_blocks, rate),
            t_last = t1

    print "\n\n  completed in %d seconds\n" % (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", default=def_blocksize,
    	help="block size (bytes, default 1M)", metavar="BYTES" )
    parser.add_option( "-c", "--cipher", dest="cipher",
    	type="string", default=def_cipher,
    	help="set ssh cipher (default: "+def_cipher+")" )
    parser.add_option( "-x", "--no-dry-run", dest="f_exec",
    	action="store_true", default=False,
    	help="really writes to output (see --dry-run option)" )
    parser.add_option( "-n", "--dry-run", dest="f_exec",
    	action="store_false",
    	help="don't writes to output (for compatibility only, this is the default)"
		+"; note that the blocks are still transferred to destination, so"
		+" you can evaluate transfer speed, they are read by destination machine"
		+" and discarded" )
    parser.add_option( "--no-sudo", dest="f_sudo",
    	action="store_false", default=True,
	help="disable use of 'sudo' command on target system; note that the remote"
		+" user needs to have write access to target device; sudo is"
		+" automatically disabled if the remote server is 'localhost'" )


    (opts, 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, opts.blocksize, opts.f_exec)
    else:
        srcdev = args[0]
        dsthost = args[1]
        if len(args) > 2:
            dstdev = args[2]
        else:
            dstdev = None

        sync(srcdev, dsthost, dstdev, opts.blocksize, opts.cipher, opts.f_exec, opts.f_sudo)


"""
Copyright 2006-2008 Justin Azoff <justin@bouncybouncy.net>
Copyright 2011 Robert Coup <robert@coup.net.nz>

== KUBiC Labs (Switzerland) patches ==

changes made by Lorenzo Canovi <kanna@kubiclabs.com>:

2023 1.4
 - removed ssh call if remote = localhost
 - disabled sudo if remote = localhost

2022 1.3
 - changed ssh cipher to 'aes192-cbc' (blowfish is deprecated)
 - added option "-c cipher"
 - added options "-n" and "-x"
 - added option "--no-sudo"

2012 1.2
 - restored original adler32 checksum, faster than sha

2012 1.1
 - changed filename from blocksync.py to blocksync
 - changed ssh remote invocation from 'python blocksync.py' to
   'blocksync'; make sure that the script is insalled in a $PATH
   directory and marked as executable

License: GPL
"""
