Cyberhack - RAID fix

cyberhack-raid.py

import os
import shutil

# function to find missing disk in block dict
def missingDisk(minDisk, maxDisk, blockDict):
    disks = []
    for disk in range(minDisk, maxDisk+1):
        try:
            blockDict[disk]
        except:
            disks.append(disk)
    return disks

# read data from our file into a bytearray
def getData(inputFile):
    return bytearray(open(inputFile, "rb").read())

# xor our two bytearrays
def xorBytearray(bytearray1, bytearray2):
    size = len(bytearray1)
    xord_byte_array = bytearray(size)
    for i in range(size):
        xord_byte_array[i] = bytearray1[i] ^ bytearray2[i]
    return xord_byte_array

tString = "disk{}_block_{:0>5}"

minDisk = 0
maxDisk = 2
found = False
filesMatch = {}
finalStuff = {}
# loop over all our blocks, make sure we get everything
# even thou there should only be 512 (0-511)
for block in range(0, 520):
    # flag to see if we found files/blocks this round
    found = False
    print("-------------")
    # loop over all our disks for this block
    for disk in range(minDisk, maxDisk+1):
        # format it!
        curBlockDisk = tString.format(disk, block)
        # check if this disk has this block
        if os.path.isfile("out/"+curBlockDisk):
            # we found something!
            found = True

            # check if we already have initialized our dict for this block
            try:
                filesMatch[block]
            except:
                filesMatch[block] = {}

            # add this disk/block to our dict for this block
            filesMatch[block][disk] = {"disk": disk, "block": block}
    # if not found, just print it and continue
    if not found:
        print("Not found in filesMatch - "+str(block))
        continue

    # find our missing blocks (spits out a list
    missing = missingDisk(minDisk, maxDisk, filesMatch[block])

    # we have three disks, and can only afford to loose one
    # if two are missing, then ERROR and exit
    if len(missing) > 1:
        print("ERROR!")
        exit()

    # could be smarter, but check if each disk is missing seperately
    # do appropiate actions for the disks and add bytearray to our finalStuff 
    # finalStuff includes ALL OUR DISK STUFF
    if 0 in missing:
        print("missing 0")
        disk1 = getData("out/"+tString.format(filesMatch[block][1]["disk"], filesMatch[block][1]["block"]))
        disk2 = getData("out/"+tString.format(filesMatch[block][2]["disk"], filesMatch[block][2]["block"]))
        disk0 = xorBytearray(disk1, disk2)
        finalStuff[block] = {
            0: disk0,
            1: disk1,
            2: disk2,
        }
    if 1 in missing:
        print("missing 1")
        disk0 = getData("out/"+tString.format(filesMatch[block][0]["disk"], filesMatch[block][0]["block"]))
        disk2 = getData("out/"+tString.format(filesMatch[block][2]["disk"], filesMatch[block][2]["block"]))
        disk1 = xorBytearray(disk0, disk2)
        finalStuff[block] = {
            0: disk0,
            1: disk1,
            2: disk2,
        }
    if 2 in missing:
        print("missing 2")
        disk0 = getData("out/"+tString.format(filesMatch[block][0]["disk"], filesMatch[block][0]["block"]))
        disk1 = getData("out/"+tString.format(filesMatch[block][1]["disk"], filesMatch[block][1]["block"]))
        disk2 = xorBytearray(disk0, disk1)
        finalStuff[block] = {
            0: disk0,
            1: disk1,
            2: disk2,
        }

    
# open handlers to all our files
disk0 = open("disk0.bin", "wb")
disk1 = open("disk1.bin", "wb")
disk2 = open("disk2.bin", "wb")
disk3 = open("disk3.bin", "wb")
for block in finalStuff:
    print("-----------")
    print(block)
    # write to disk 0..2
    disk0.write(finalStuff[block][0])
    disk1.write(finalStuff[block][1])
    disk2.write(finalStuff[block][2])
    # write d00b00 and d01b00 first, then
    # write d00b01 and d01b01
    disk3.write(finalStuff[block][0])
    disk3.write(finalStuff[block][1])
disk0.close()
disk1.close()
disk2.close()