#!/usr/bin/python3
# -*- coding: utf-8 -*-

# Script for reverting ZFS changes by destroying uberblocks
# Author: Martin Vool
# E-mail: mardicas@gmail.com
# Version: 0.1
# Date: 16 November 2009

import time
import subprocess
import sys
import os
from pathlib import Path
import click


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def verify(thing):
    if not thing:
        raise ValueError(thing)


def get_block_device_size(device):
    assert Path(device).is_block_device()
    fd = os.open(device, os.O_RDONLY)
    try:
        return os.lseek(fd, 0, os.SEEK_END)
    finally:
        os.close(fd)


@click.command()
@click.argument("device", type=click.Path(exists=True, dir_okay=False))
@click.option("--bs", type=int, default=512, help="blocksize")  # todo sanity check
@click.option("--tb", type=int, required=True, help="total device size in blocks")
@click.option("--verbose", is_flag=True)
def main(device, bs, tb, verbose):
    device_size_bytes = get_block_device_size(device)
    verify(device_size_bytes % bs == 0)
    verify(device_size_bytes / bs == tb)

    eprint("device:", device)
    eprint("bs:", bs)
    eprint("tb:", bs)

    # make solaris use gnu grep.
    if os.uname()[0] == 'SunOS':
        grep_cmd = 'ggrep'
    else:
        grep_cmd = 'grep'

    # to format program output
    def formatstd(inp):
        verify(isinstance(inp, bytes))
        inp = inp.split(b'\n')
        ret = []
        for line in inp:
            columns = line.split(b' ')
            nc = []
            for c in columns:
                if c != b'':
                    nc.append(c)
            ret.append(nc)
        return ret

    # read blocks from beginning(64mb)
    a_count = (256 * bs)
    # read blocks from end (64mb)
    l_skip = tb - (256 * bs)
    if verbose:
        eprint("a_count:", a_count)
        eprint("l_skip:", l_skip)

    eprint('Total of %s blocks' % tb)
    eprint('Reading from the beginning to %s blocks' % a_count)
    eprint('Reading from %s blocks to the end' % l_skip)

    # get the uberblocks from the beginning and end
    command = 'sync && dd bs=%s if=%s count=%s | od -A x -x | %s -A 2 "b10c 00ba" | %s -v "\-\-"'
    yberblocks_a = \
        formatstd(
            subprocess.Popen(
                command % (bs, device, a_count, grep_cmd, grep_cmd), shell=True, stdout=subprocess.PIPE).communicate()[0])

    command = 'sync && dd bs=%s if=%s skip=%s | od -A x -x | %s -A 2 "b10c 00ba" | %s -v "\-\-"'
    yberblocks_l = \
        formatstd(
            subprocess.Popen(
                command % (bs, device, l_skip, grep_cmd, grep_cmd), shell=True, stdout=subprocess.PIPE).communicate()[0])

    if verbose:
        eprint("yberblocks_a:", yberblocks_a)
        eprint("yberblocks_l:", yberblocks_l)

    yberblocks = []

    for p in yberblocks_a:
        if len(p) > 0:
            # format the hex address to decmal so dd would eat it.
            p[0] = int(int(p[0], 16) / bs)
            yberblocks.append(p)

    for p in yberblocks_l:
        if len(p) > 0:
            # format the hex address to decmal so dd would eat it and add the skipped part.
            p[0] = int((int(p[0], 16) / bs) + int(l_skip))  # we add until the place we skipped so the adresses will mach.
            yberblocks.append(p)

    if verbose:
        eprint("yberblocks:")
        for y in yberblocks:
            eprint(y)

    print('----')
    # here will be kept the output that you will see later(TXG, timestamp and the adresses, should be 4, might be less)
    koik = {}
    i = 0
    for p in yberblocks:
        if len(p) > 0:
            if i == 0:  # the first output line
                address = p[0]
            elif i == 1:  # second output line
                # this is the output of od that is in hex and needs to be reversed
                txg = int(p[4] + p[3] + p[2] + p[1], 16)
            elif i == 2:  # third output line
                timestamp = int(p[4] + p[3] + p[2] + p[1], 16)
                try:
                    aeg = time.strftime("%d %b %Y %H:%M:%S", time.localtime(timestamp))
                except Exception as e:
                    print(e)
                    aeg = 'none'
                if txg in koik:
                    koik[txg]['addresses'].append(address)
                else:
                    koik[txg] = {
                        'txg': txg,
                        'timestamp': timestamp,
                        'htime': aeg,
                        'addresses': [address]
                    }
            if i == 2:
                i = 0
            else:
                i += 1

        keys = list(koik.keys())
        keys.sort()

        if verbose:
            eprint("keys:", keys)

    while True:
        keys = list(koik.keys())
        keys.sort()
        print('TXG\tTIME\tTIMESTAMP\tBLOCK ADDRESSES')
        for k in keys:
            print('%s\t%s\t%s\t%s' % (k, koik[k]['htime'], koik[k]['timestamp'], koik[k]['addresses']))
        try:
            save_txg = int(eval(input('What is the last TXG you wish to keep? (CTRL-c if done)\n')))
            verify(save_txg in keys)
            keys = list(koik.keys())
            keys.sort()
            for k in keys:
                if k > save_txg:
                    for adress in koik[k]['addresses']:
                        # wrtie zeroes to the unwanted uberblocks
                        command = 'dd bs=%s if=/dev/zero of=%s seek=%s count=1 conv=notrunc' % (bs, device, adress)
                        eprint("command:", command)
                        fmt = formatstd(subprocess.Popen(command, shell=True, stdout=subprocess.PIPE).communicate()[0])
                    del koik[k]
            # sync changes to disc!
            sync = formatstd(subprocess.Popen('sync', shell=True, stdout=subprocess.PIPE).communicate()[0])
        except Exception as e:
            print(e)
            print('')
            break

if __name__ == '__main__':
    main()