#! /usr/bin/env python
#   @file       fat32bstool.py
#   @copyright  Copyright 2009 David Elliott. All rights reserved.
#   @abstract   Updates FAT32 boot blocks on a disk device
#   @discussion
#   Right now it only accepts (and requires) -B option to specify the file
#   containing the new bootstrap code.

import sys
import struct
import getopt

def isJumpValid(buf):
    return (buf[0] == "\xeb" and buf[2] == "\x90") or (buf[0] == "\xe8")

def isFSInfoValid(buf):
    if buf[0x200:0x204] != "RRaA":
        return False
    if buf[0x3e4:0x3e8] != "rrAa":
        return False
    if buf[0x3fc:0x400] != "\x00\x00\x55\xaa":
        return False
    return True

def isReservedSectorsSufficient(buf):
    BPB_BytsPerSec, BPB_SecPerClus, BPB_RsvdSecCnt = struct.unpack("<HBH", buf[11:16])
    return BPB_BytsPerSec == 512 and BPB_RsvdSecCnt >= 3


def main(argv):
    # -B bootcode was chosen to match newfs_msdos
    # -o hidden sets the number of hidden sectors.
    opts, args = getopt.getopt(argv[1:], "B:u:h:o:")

    codeFN = None
    hiddenSectors = None
    BPB_SecPerTrk = None
    BPB_NumHeads = None
    for opt in opts:
        if opt[0] == "-B":
            codeFN = opt[1]
        if opt[0] == "-u":
            BPB_SecPerTrk = int(opt[1],0)
        if opt[0] == "-h":
            BPB_NumHeads = int(opt[1],0)
        if opt[0] == "-o":
            hiddenSectors = int(opt[1],0)

    outF = None

    if len(args) < 1 or len(args) > 2:
        print >> sys.stderr, "Usage %s -B bootcode diskDevice [output]" % argv[0]
        print >> sys.stderr, "If output is not specified, updated boot blocks are written back to diskDevice"
        return 1

    deviceFN = args[0]

    if len(args) == 2:
        if args[1] == "-":
            outF = sys.stdout
        else:
            outF = open(args[1], "wb")

    deviceF = open(deviceFN, "rb")
    deviceBuf = deviceF.read(512 * 3)
    deviceF.close()
    deviceF = None
    if len(deviceBuf) != 512 * 3:
        print >> sys.stderr, "Could not read 3 512-byte sectors from device %s" % deviceFN
        return 1
    codeBuf = None
    if codeFN != None:
        codeF = open(codeFN, "rb")
        codeBuf = codeF.read(512 * 3)
        codeF.close()
        codeF = None
    else:
        codeBuf = deviceBuf
    if len(codeBuf) != 512 * 3:
        print >> sys.stderr, "Could not read 3 512-byte sectors from code file %s" % codeFN
        return 1
    if not isJumpValid(codeBuf):
        print >> sys.stderr, "Code does not begin with a jump instruction"
        return 1
    if not isJumpValid(deviceBuf):
        print >> sys.stderr, "Device does not begin with a jump instruction"
        return 1
    if not isReservedSectorsSufficient(deviceBuf):
        print >> sys.stderr, "Device must have at least 3 reserved 512-byte sectors"
    if not isFSInfoValid(codeBuf):
        print >> sys.stderr, "Code does not have valid FSInfo signatures"
        return 1
    if not isFSInfoValid(deviceBuf):
        print >> sys.stderr, "Device does not have valid FSInfo signatures"
        #return 1

    hiddenSectorBuf = None
    if hiddenSectors != None:
        hiddenSectorBuf = struct.pack("<L", hiddenSectors)
    else:
        hiddenSectorBuf = deviceBuf[0x1c:0x20]
    sptBuf = struct.pack("<H", BPB_SecPerTrk) if (BPB_SecPerTrk != None) else deviceBuf[0x18:0x1a]
    headsBuf = struct.pack("<H", BPB_NumHeads) if (BPB_NumHeads != None) else deviceBuf[0x1a:0x1c]

    outbuf = codeBuf[0:3] + deviceBuf[3:0x18] + sptBuf + headsBuf + hiddenSectorBuf + deviceBuf[0x20:0x5a] + codeBuf[0x5a:0x200] + deviceBuf[0x200:0x400] + codeBuf[0x400:0x600]

    if outF != None:
        outF.write(outbuf)
    else:
        # Reopen the device in read/write mode w/o truncate
        deviceF = open(deviceFN, "r+")
        byteswritten = deviceF.write(outbuf)
        print >> sys.stderr, "Wrote %d bytes back to %s" % (len(outbuf), deviceFN)
        deviceF.close()

if __name__ == "__main__":
    main(sys.argv)