#!/usr/bin/env python3


"""
Synchronize block devices. Blocksync does for block devices what rsync does for files.

Blocksync does exactly what the name suggests. It copies changes from the source block device to
the destination block device. To do this both source and destination are read chunk by chunk and
then block checksums are compared. If a checksum differs, the chunk is transferred. Both source and
destination can be either local or remote and either a chunk device or a file. The only dependency
is Python3 on both source and destination machines.

Written 2020 by Maarten Tromp <maarten@geekabit.nl>


Website
-------
https://www.geekabit.nl/projects/blocksync/


Usage
-----
Synchronize two local devices:
blocksync --verbose /dev/sdb /dev/sdc

Synchronize local block device and file:
blocksync --verbose /dev/sdb sdb.img

Push changes to a remote block device:
blocksync --verbose /dev/sdb remote-host.org:/dev/sdc

Pull changes from a remote block device:
blocksync --verbose user@remote-host.org:/dev/sdc sdc.img


Inspiration
-----------
This tool is inspired by another project named blocksync. This code is written from scratch, using
roughly the same approach as the other project.

blocksync as found at: https://github.com/shodanshok/blocksync
	Justin Azoff <justin@bouncybouncy.net>
	Robert Coup <robert@coup.net.nz>
	Gionatan Danti <g.danti@assyoma.it>


License
-------
CC0 (no copyright, public domain)
The person who associated a work with this deed has dedicated the work to the public domain by
waiving all of his or her rights to the work worldwide under copyright law, including all related
and neighbouring rights, to the extent allowed by law. You can copy, modify, distribute and perform
the work, even for commercial purposes, all without asking permission.
https://creativecommons.org/share-your-work/public-domain/cc0/


Bugs / limitations
------------------
- Source and destination size need to be a multiple of the chunk size. With 1MB chunk size, this
  usually is the case.
- While being more efficient than copying an entire block device, blocksync still has to read the
  entire source and destination to determine changes.
"""


# imports
import sys
import argparse
import subprocess
import inspect
import datetime


def worker():
	"""
	Interface with block device or file, communicates with control process.

	This is the worker process that will run on both source and destination host. It reads / writes
	the target file / block device and connects to control process via pipes.


	Note
	----
	The function below will be started as subprocess and should be self-contained.
	"""

	import sys
	import hashlib

	chunk_size = 1024 * 1024 # bytes

	# send ready
	sys.stdout.buffer.write(("ready\n").encode())
	sys.stdout.buffer.flush()

	# receive parameters
	# role: string: "src" or "dst"
	role = sys.stdin.buffer.readline().decode().strip()
	# file_name: string: file or block device to open
	file_name = sys.stdin.buffer.readline().decode().strip()

	if role == "src":
		open_mode = "rb" # read-only
	elif role == "dst":
		open_mode = "rb+" # read and update
	else:
		sys.stderr.buffer.write((f"ERROR unknown role: {role}\n").encode())
		sys.exit(1)

	# if file_name does not exist of has incorrect permissions, it will crash here
	# which is no problem since it will be gracefully handled by the control function
	file_handle = open(file_name, open_mode)

	# send device size (string)
	file_handle.seek(0, 2)
	file_size = file_handle.tell()
	sys.stdout.buffer.write((f"{file_size}\n").encode())
	sys.stdout.buffer.flush()

	chunk_number = 0
	while True:
		# read chunk
		file_handle.seek(chunk_number * chunk_size)
		chunk_data = file_handle.read(chunk_size)

		# send checksum (binary)
		# no need for a very secure hash, small digest size limits overhead.
		checksum = hashlib.sha1(chunk_data).digest()
		sys.stdout.buffer.write(checksum)
		sys.stdout.buffer.flush()

		# receive next command
		cmd = sys.stdin.buffer.readline().decode().strip()
		if cmd == "send" and role == "src":
			# send entire chunk (binary)
			sys.stdout.buffer.write(chunk_data)
			sys.stdout.buffer.flush()
		# TODO: add commands "send-partial" and "receive-partial" to transfer incomplete last chunk. Add
		# chunk length as parameter.
		#if cmd == "send-partial" and role == "src":
		#	# send partial chunk (binary)
		#	sys.stdout.buffer.write(chunk_length)
		#	sys.stdout.buffer.write(chunk_data)
		#	sys.stdout.buffer.flush()
		elif cmd == "receive" and role == "dst":
			# receive entire chunk (binary)
			chunk_data = sys.stdin.buffer.read(chunk_size)
			# write received chunk
			file_handle.seek(chunk_number * chunk_size)
			file_handle.write(chunk_data)
		#elif cmd == "receive-partial" and role == "dst":
		#	# receive partial chunk (binary)
		#	chunk_size = sys.stdin.buffer.read(??)
		#	chunk_data = sys.stdin.buffer.read(chunk_size)
		#	# write received chunk
		#	file_handle.seek(chunk_number * chunk_size)
		#	file_handle.write(chunk_data)
		elif cmd == "next":
			pass
		else:
			sys.stderr.buffer.write((f"ERROR unknown command: {cmd}\n").encode())
			sys.exit(1)

		# advance to next chunk
		chunk_number += 1


def control(config):
	"""
	Controls worker subprocesses.

	This is the main function. Subprocesses are spawned for reading and writing source and destination
	block devices.
	"""

	chunk_size = 1024 * 1024

	# add boilerplate around worker process so it can be started as Python command-line parameter
	sub_cmd = "/usr/bin/env python3 -c '''" # Make sure these triple quotes are not used within the function.
	sub_cmd += inspect.getsource(worker) # copy literal function source
	sub_cmd += "worker()" # call worker process
	sub_cmd += "'''"

	# setup subprocesses
	src_file_name = config.src
	if ":" in src_file_name:
		# remote
		src_host, src_file_name = src_file_name.split(":")
		src_cmd = ("ssh", "--", src_host, sub_cmd)
		src_shell = False
	else:
		# local
		src_host = ""
		src_cmd = sub_cmd
		src_shell = True

	dst_file_name = config.dst
	if ":" in dst_file_name:
		dst_host, dst_file_name = dst_file_name.split(":")
		dst_cmd = ("ssh", "--", dst_host, sub_cmd)
		dst_shell = False
	else:
		dst_host = ""
		dst_cmd = sub_cmd
		dst_shell = True

	# spawn subprocesses
	src_p = subprocess.Popen(args = src_cmd, stdin = subprocess.PIPE, stdout = subprocess.PIPE,
			stderr = subprocess.PIPE, shell = src_shell)
	dst_p = subprocess.Popen(args = dst_cmd, stdin = subprocess.PIPE, stdout = subprocess.PIPE,
			stderr = subprocess.PIPE, shell = dst_shell)
	# Do not use unbuffered pipes, it leads to all kinds of problems when transfering bigger chunks of data.

	# wait for subprocesses to report ready (string)
	src_stdout = src_p.stdout.readline().decode().strip()
	if src_stdout != "ready":
		sys.stderr.write(src_p.stderr.read().decode())
		sys.stderr.write(src_p.stderr)
		sys.exit(1)
	if config.verbose >= 3: # debug
		print("src subprocess ready")

	dst_stdout = dst_p.stdout.readline().decode().strip()
	if dst_stdout != "ready":
		sys.stderr.write(dst_p.stderr.read().decode())
		sys.stderr.write(dst_p.stderr)
		sys.exit(1)
	if config.verbose >= 3: # debug
		print("dst subprocess ready")

	# send parameters (string)
	src_p.stdin.write((f"src\n{src_file_name}\n").encode())
	src_p.stdin.flush()

	dst_p.stdin.write((f"dst\n{dst_file_name}\n").encode())
	dst_p.stdin.flush()

	# receive sizes (string)
	src_size = src_p.stdout.readline().decode().strip()
	if config.verbose >= 3: # debug
		print(f"src size: {src_size}")
	if not src_size:
		sys.stderr.write(f"ERROR unable to open: {config.src}\n")
		sys.exit(1)
	src_size = int(src_size)

	dst_size = dst_p.stdout.readline().decode().strip()
	if config.verbose >= 3: # debug
		print(f"dst size: {dst_size}")
	if not dst_size:
		sys.stderr.write(f"ERROR unable to open: {config.dst}\n")
		sys.exit(1)
	dst_size = int(dst_size)

	if src_size > dst_size:
		sys.stderr.write("ERROR source is bigger than destination\n")
		sys.exit(1)
	if src_size % chunk_size != 0:
		sys.stderr.write("ERROR source size is not multiple of chunk size\n")
		sys.exit(1)

	chunks_total = src_size / chunk_size
	if config.verbose >= 3: # debug
		print(f"src chunks total: {chunks_total}")
	chunks_done = 0
	chunks_same = 0
	chunks_different = 0
	datetime_start = datetime.datetime.now()

	# start comparing chunks
	stop = False
	while not stop:
		# show status
		pct_done = int(100 * chunks_done / chunks_total)
		time_passed = str(datetime.datetime.now() - datetime_start)
		if config.verbose == 2: # verbose only
			print(f"\r{pct_done}%, chunk: {chunks_done}/{chunks_total}, same: {chunks_same}, different: {chunks_different}, time: {time_passed}", end = "")
			# "\r" resets cursor to start of line, end="" doesn't end line in "\n", so we keep overwriting the same line
		elif config.verbose >= 3: # debug
			print(f"{pct_done}%, chunk: {chunks_done}/{chunks_total}, same: {chunks_same}, different: {chunks_different}, time: {time_passed}")
			# same thing but without line overwrite magic

		# receive checksums (binary)
		src_checksum = src_p.stdout.read(20) # sha1 digest size = 20 bytes
		dst_checksum = dst_p.stdout.read(20)
		if src_checksum == dst_checksum:
			chunks_same += 1
			if config.verbose >= 3: # debug
				print("chunks same")
			src_p.stdin.write(("next\n").encode())
			src_p.stdin.flush()
			dst_p.stdin.write(("next\n").encode())
			dst_p.stdin.flush()
		else:
			chunks_different += 1
			if config.verbose >= 3: # debug
				print("chunks different")
			src_p.stdin.write(("send\n").encode())
			src_p.stdin.flush()
			dst_p.stdin.write(("receive\n").encode())
			dst_p.stdin.flush()
			# receive chunk (binary)
			chunk_data = src_p.stdout.read(chunk_size)
			dst_p.stdin.write(chunk_data)
			src_p.stdin.flush()
		chunks_done += 1

		if chunks_done == chunks_total:
			stop = True

	if config.verbose >= 2: # verbose
		print() # newline
	if config.verbose >= 1: # info
		size_mb = src_size / (1024 * 1024)
		copied_mb = chunks_different * chunk_size / (1024 * 1024)
		seconds_passed = (datetime.datetime.now() - datetime_start).total_seconds()
		avg_speed_mbps = src_size / seconds_passed / (1024 * 1024)
		print(f"Copied {copied_mb} MB / {size_mb} MB in {time_passed}, speed: {avg_speed_mbps:.2f} MB/s")


if __name__ == "__main__":
	# parse command-line parameters
	parser = argparse.ArgumentParser(description = "%prog [OPTION]... [[USER@]HOST:]SRC [[USER@]HOST:]DST")
	parser.add_argument("src", help = "Source file or block device. This can be either local or remote.")
	parser.add_argument("dst", help = "Destination file or block device. This can be either local or remote.")
	parser.add_argument("-v", "--verbose", action = "count", default = 1, help = "Increase verbosity")
	parser.add_argument("-q", "--quiet", action = "store_true", default = False, help="Suppress all output")
	args = parser.parse_args()

	if args.quiet:
		args.verbose = 0
	# verbosity levels:
	# 0 = quiet, no output except errors
	# 1 = info (default), only print "completed in ..." message
	# 2 = verbose, include progress output
	# 3 = debug, include variable dumps

	if args.verbose >= 3: # debug
		print(f"src device: {args.src}")
		print(f"dst device: {args.dst}")
		print(f"verbose: {args.verbose}")

	# start control function
	control(args)