#! /usr/bin/env python

import os, subprocess, threading

dd_command =   "dd bs=1K count=%i status=noxfer iflag=direct conv=noerror skip=%i if=%s"

device = "/dev/sdd1"
output_file = open("mirror.img", 'wb')

timeout  = 5

def check_device(device):
    while not os.path.exists(device):
        raw_input("Jiggle the drive and press Enter.")

def call_stdout(cmd, timeout):
    pipe = subprocess.PIPE
    cmd_kill = threading.Timer(timeout, lambda: subp.kill())
    cmd_kill.start()
    subp = subprocess.Popen(cmd, shell=True, stdin=None, stdout=pipe, stderr=pipe)
    raw_data = ''
    raw_data = subp.communicate()[0]
    cmd_kill.cancel()
    if subp.returncode != 0:
        return subp.returncode
    return raw_data

def main():
    # pos and size in KB (1024)
    todo = [(pos, 4*1024) for pos in range(0, 9000*1024, 4*1024)]
    badblocks = 0
    while todo:
        check_device(device)
        pos,size = todo.pop(0)
        print '%iM for %ik' % (pos // 1024, size),
        raw_data = call_stdout(dd_command % (size, pos, device), timeout)
        if raw_data == 1:  # EOF?
            print 'EOF'
            output_file.close()
            return
        if raw_data == -9 and size == 1:
            output_file.write('\00' * 1024)
            print 'badblocks + 1 =', badblocks + 1 
            badblocks += 1
            continue
        if type(raw_data) == str and len(raw_data) == size*1024:
            output_file.write(raw_data)
            print
        elif type(raw_data) == str and len(raw_data) == 0:
            raw_input("\nJiggle the drive and press Enter.")
            check_device(device)
            todo.insert(0, (pos, size))
        else:
            raw1 = call_stdout(dd_command % (size//2, pos, device), timeout)
            raw2 = call_stdout(dd_command % (size//2, pos+size//2, device), timeout)
            if raw1 == raw2 == -9:
                output_file.write('\00' * 1024 * size)
                print 'badblocks +', size, '=', badblocks + size
                badblocks += size
            else:
                print 'semifail'
                todo.insert(0, (pos+size//2, size//2))
                todo.insert(0, (pos,         size//2))

main()