import struct, os, sys
import getopt
import time
import hashlib
import base64
import ctypes


# Constants
SECTOR_SIZE            = 512
FILE_ENTRY_SIZE        = 32
FAT_FILE_NAME_LENGTH   = 11
ATTRIBUTE_READONLY     = 0x01
ATTRIBUTE_HIDDEN       = 0x02
ATTRIBUTE_SYSTEM       = 0x04
ATTRIBUTE_VOL_LABEL    = 0x08
ATTRIBUTE_DIR          = 0x10
ATTRIBUTE_ARCHIVE      = 0x20
FAT16_ENTRY_SIZE       = 2
FAT32_ENTRY_SIZE       = 4
DIR_ENT_LENGTH         = 32

# Structs
fat_common_header = struct.Struct('<3s8sHBHBHHBHHHII')
fat_16_header = struct.Struct('<HBI11s8s448s2s')
fat_dirent = struct.Struct('<8s3sBBBHHHHHHHI')
fat_32_header = struct.Struct('<IHHIHH12sBBBI11s8s420s2s')
fat_32_fsi = struct.Struct('<4s480s4sII12s2s2s')
long_name_dirent = struct.Struct('<B10sBBB12sH4s')

# Info dictionaries
fat_header_info = {}
fat_32_fsi_info = {}

# Global properties
fat_container_file = None
fat_type = None
current_cluster = None
chain_start = None


def round_up(dividend, divisor):
    result = dividend / divisor
    if dividend % divisor != 0:
        result = result + 1
    return result


def PrintBigError(sz):
    print "\t _________________ ___________ "
    print "\t|  ___| ___ \\ ___ \\  _  | ___ \\"
    print "\t| |__ | |_/ / |_/ / | | | |_/ /"
    print "\t|  __||    /|    /| | | |    / "
    print "\t| |___| |\\ \\| |\\ \\\\ \\_/ / |\\ \\ "
    print "\t\\____/\\_| \\_\\_| \\_|\\___/\\_| \\_|\n"

    if len(sz) > 0:
        print sz
        sys.exit(1)


def read_header():
    global fat_header_info, fat_type, fat_32_fsi_info

    try:
        fp = open(fat_container_file, 'rb')
    except:
        PrintBigError("Can't open '%s' in binary read mode" % fat_container_file)

    try:
        fp.seek(0)
    except:
        PrintBigError("Can't seek to 0 in '%s'" % fat_container_file)

    try:
        data = fp.read(fat_common_header.size)
    except:
        PrintBigError("Can't read %d bytes from '%s'" % (fat_common_header.size, fat_container_file))

    try:
        fat_data = fp.read(SECTOR_SIZE - fat_common_header.size)
    except:
        PrintBigError("Can't read %d bytes from '%s'" % (SECTOR_SIZE - fat_common_header.size, fat_container_file))

    try:
        fp.close()
    except:
        PrintBigError("Can't close '%s'" % (fat_container_file))

    all_bytes = fat_common_header.unpack(data)
    fat_header_info['jump_code_nop'] = all_bytes[0]
    fat_header_info['oem_name'] = all_bytes[1]
    fat_header_info['bytes_per_sec'] = all_bytes[2]
    if SECTOR_SIZE != fat_header_info['bytes_per_sec']:
        PrintBigError("Sector size mismatch. Header %d Specified %d" % (fat_header_info['bytes_per_sec'], SECTOR_SIZE))
    fat_header_info['sec_per_cluster'] = all_bytes[3]
    fat_header_info['reserved_sec'] = all_bytes[4]
    fat_header_info['num_fat_copies'] = all_bytes[5]
    fat_header_info['max_root_dir_ent'] = all_bytes[6]
    fat_header_info['num_sec_small_32mb'] = all_bytes[7]
    fat_header_info['media_desc'] = all_bytes[8]
    fat_header_info['sec_per_fat'] = all_bytes[9]
    fat_header_info['sec_per_track'] = all_bytes[10]
    fat_header_info['num_heads'] = all_bytes[11]
    fat_header_info['num_hidden_sec'] = all_bytes[12]
    fat_header_info['num_sec'] = all_bytes[13]
    print "Container size is %d bytes" % (int(fat_header_info['num_sec']) * SECTOR_SIZE)

    if fat_data[18:18 + 8].strip().lower() == 'FAT16'.lower():
        fat_type = 'fat16'
        fat_data = fat_data[:fat_16_header.size]
        all_bytes = fat_16_header.unpack(fat_data)
        fat_header_info['logical_drive_num'] = all_bytes[0]
        fat_header_info['ext_sign'] = all_bytes[1]
        fat_header_info['ser_num'] = all_bytes[2]
        fat_header_info['vol_name'] = all_bytes[3]
        fat_header_info['fat_name'] = all_bytes[4]
        # fat_header_info['exec_code'] = all_bytes[5]
        fat_header_info['exec_marker'] = all_bytes[6]
    else:
        fat_type = 'fat32'
        fat_data = fat_data[:fat_32_header.size]
        all_bytes = fat_32_header.unpack(fat_data)
        fat_header_info['old_sec_per_fat'] = fat_header_info['sec_per_fat']
        fat_header_info['sec_per_fat'] = all_bytes[0]
        fat_header_info['flags'] = all_bytes[1]
        fat_header_info['ver_num'] = all_bytes[2]
        fat_header_info['root_dir_start_cluster'] = all_bytes[3]
        fat_header_info['fsi_start_num'] = all_bytes[4]
        fat_header_info['bkp_boot_sec_num'] = all_bytes[5]
        # fat_header_info['reserved'] = all_bytes[6]
        fat_header_info['logical_drive_num'] = all_bytes[7]
        # fat_header_info['unused'] = all_bytes[8]
        fat_header_info['ext_sign'] = all_bytes[9]
        fat_header_info['ser_num'] = all_bytes[10]
        fat_header_info['vol_name'] = all_bytes[11]
        fat_header_info['fat_name'] = all_bytes[12]
        # fat_header_info['exec_code'] = all_bytes[13]
        fat_header_info['exec_marker'] = all_bytes[14]

        try:
            fp = open(fat_container_file, 'rb')
        except:
            PrintBigError("Can't open '%s' in binary read mode" % (fat_container_file))

        try:
            fp.seek(fat_header_info['fsi_start_num'] * SECTOR_SIZE)
        except:
            PrintBigError("Can't seek to location %d in '%s'" % (
            fat_header_info['fsi_start_num'] * SECTOR_SIZE, fat_container_file))

        all_bytes = fat_32_fsi.unpack(fp.read(SECTOR_SIZE))

        try:
            fp.close()
        except:
            PrintBigError("Can't close '%s'" % (fat_container_file))

        fat_32_fsi_info['fsi_sign_1'] = all_bytes[0]
        # fat_32_fsi_info['reserved'] = all_bytes[1]
        fat_32_fsi_info['fsi_sign_2'] = all_bytes[2]
        fat_32_fsi_info['num_free_clusters'] = all_bytes[3]
        fat_32_fsi_info['recent_alloc_cluster'] = all_bytes[4]
        # fat_32_fsi_info['reserved'] = all_bytes[5]
        # fat_32_fsi_info['reserved'] = all_bytes[6]
        fat_32_fsi_info['fsi_marker'] = all_bytes[7]

def ChkSum(short_name):
    dot = short_name.find('.')
    if dot == -1:
        while len(short_name) != 11:
            short_name += " "
    else:
        file = short_name[:dot]
        ext = short_name[dot + 1: dot + 1 + 3]
        while len(file) != 8:
            file += " "
        while len(ext) != 3:
            ext += " "
        short_name = file + ext

    sum = 0
    for c in short_name:
        sum = ((((sum & 1) << 7) | ((sum & 0xfe) >> 1)) + ord(c)) & 0xff
        #sum = 0x80 if (sum & 1) else 0 + (sum >> 1) + ord(c)
        #sum = sum & 0xff
    return sum

def Split_File_Name(file_name, count):
    split = []
    i = 0
    offset = 0
    flag = 0
    while i < count:
        name1 = file_name[offset+0:offset+5]
        if len(name1) != 5:
            if flag == 0:
                name1 += u'\x00'
                flag = 1
        name2 = file_name[offset+5:offset+11]
        if len(name2) != 6:
            if flag == 0:
                name2 += u'\x00'
                flag = 1
        name3 = file_name[offset+11:offset+13]
        if len(name3) != 2:
            if flag == 0:
                name3 += u'\x00'
                flag = 1
        split.append(name1)
        split.append(name2)
        split.append(name3)
        i = i + 1
        offset = i * 13
    return split

def GenerateLongEntry(file_name, short_name, count):
    short_name = short_name.replace('.', '')
    chksum = ChkSum(short_name)
    long_entries = []
    name_entries = Split_File_Name(file_name, count)
    i = 0
    offset = 0
    while i < count:
        ent = None
        if i == count - 1:
            # Last entry needs to be masked
            ent = long_name_dirent.pack((i + 1) | 0x40, name_entries[offset].encode('utf-16le') + b'\xff\xff' * (5-len(name_entries[offset])), 0x0f, 0, chksum, name_entries[offset + 1].encode('utf-16le') + b'\xff\xff' * (6-len(name_entries[offset + 1])), 0, name_entries[offset + 2].encode('utf-16le') + b'\xff\xff' * (2-len(name_entries[offset + 2])))
        else:
            ent = long_name_dirent.pack((i + 1), name_entries[offset].encode('utf-16le'), 0x0f, 0, chksum, name_entries[offset + 1].encode('utf-16le'), 0, name_entries[offset + 2].encode('utf-16le'))
        long_entries.append(ent)
        i = i + 1
        offset = i * 3
    return long_entries

def generate_short_name(name):
    name = name.encode("utf-8")
    hash = hashlib.md5(name)
    hex = hash.digest()
    encoded = base64.b32encode(hex)
    # Some OS will use the first 6 characters of long file names and add ~1 for the first file,
    # ~2 for the second and so on up to ~9. After running out of options it uses a hash and adds ~2.
    # This code will only use a hash for all short file names entries,
    # uses ~2 to make it similar to existing implementations.
    val = encoded[:6] + "~2"
    dot_index = name.rfind(".")
    # If there is an extension
    if dot_index != -1:
        ext = name[dot_index + 1: dot_index + 1 + 3]
        return "%s.%s" % (val.encode("utf-8"), ext.upper().encode("utf-8"))
    else:
        return val

def clean_short_dos_name(file_name):
    if "." in file_name:
        file_part = file_name[:min(8, file_name.index("."))]
        ext_part = file_name[file_name.rindex(".") + 1:file_name.rindex(".") + 1 + 3]
        return "%s.%s" % (file_part.upper(), ext_part.upper())
    else:
        return file_name[:8].upper()


def get_cluster_sector(cluster_num):
    if fat_type.lower() == "fat16":
        return fat_header_info['reserved_sec'] + \
            (fat_header_info['num_fat_copies'] * fat_header_info['sec_per_fat']) + \
            (round_up((fat_header_info['max_root_dir_ent'] * 32), fat_header_info['bytes_per_sec'])) + \
            ((cluster_num - 2) * fat_header_info['sec_per_cluster'])
    if fat_type.lower() == "fat32":
        return fat_header_info['reserved_sec'] + \
            (fat_header_info['num_fat_copies'] * fat_header_info['sec_per_fat']) + \
            ((cluster_num - 2) * fat_header_info['sec_per_cluster'])


def read_cluster(cluster_num):
    if get_cluster_sector(cluster_num) * fat_header_info['bytes_per_sec'] >= os.fstat(fat_container_file_fp.fileno()).st_size:
        write_cluster(cluster_num, '')
    fat_container_file_fp.seek(get_cluster_sector(cluster_num) * fat_header_info['bytes_per_sec'])
    return fat_container_file_fp.read(fat_header_info['sec_per_cluster'] * fat_header_info['bytes_per_sec'])


def write_cluster(cluster_num, data):
    padding_data_len = (get_cluster_sector(cluster_num) * fat_header_info['bytes_per_sec']) - os.fstat(
        fat_container_file_fp.fileno()).st_size
    if padding_data_len >= 0:
        fat_container_file_fp.seek(0, os.SEEK_END)
        fat_container_file_fp.write('\x00' * padding_data_len)
        fat_container_file_fp.seek(0, os.SEEK_END)
    else:
        fat_container_file_fp.seek(get_cluster_sector(cluster_num) * fat_header_info['bytes_per_sec'])
    data = data + '\x00' * ((fat_header_info['sec_per_cluster'] * fat_header_info['bytes_per_sec']) - len(data))
    fat_container_file_fp.write(data)
    # Check if overflow occurred after writing
    if os.fstat(fat_container_file_fp.fileno()).st_size > int(fat_header_info['num_sec']) * SECTOR_SIZE:
        temp = fat_container_file.index(".")
        file_part = fat_container_file[:temp] + "_overflow"
        ext_part = fat_container_file[temp:]
        os.rename(fat_container_file, file_part + ext_part)
        PrintBigError("Image overflow: Files added to %s are larger than the FAT image size" % fat_container_file)
        sys.exit(1)


def get_cluster_table():
    fat_container_file_fp.seek(fat_header_info['reserved_sec'] * fat_header_info['bytes_per_sec'])
    cluster_table = fat_container_file_fp.read(fat_header_info['sec_per_fat'] * fat_header_info['bytes_per_sec'])
    if fat_type.lower() == "fat16":
        num_entries = round_up(len(cluster_table), FAT16_ENTRY_SIZE)
        temp = struct.unpack("<%dh" % num_entries, cluster_table)
        cluster_table = [x & 0xffff for x in temp]
    if fat_type.lower() == "fat32":
        num_entries = round_up(len(cluster_table), FAT32_ENTRY_SIZE)
        temp = struct.unpack("<%di" % num_entries, cluster_table)
        cluster_table = [x & 0xffffffff for x in temp]
    return cluster_table


def find_free_cluster():
    cluster_table = get_cluster_table()
    if 0 in cluster_table:
        temp_free = cluster_table.index(0)
        write_cluster(temp_free, '')
        return temp_free
        # return cluster_table.index(0) == 2 and 3 or cluster_table.index(0)
    return None


def update_fat_cluster(cluster_num, next_cluster=None):
    offset = None
    value = None
    if fat_type.lower() == "fat16":
        offset = cluster_num * FAT16_ENTRY_SIZE
        if next_cluster is None:
            next_cluster = 0xffff
        value = struct.pack("<H", 0xffff & next_cluster)
    if fat_type.lower() == "fat32":
        offset = cluster_num * FAT32_ENTRY_SIZE
        if next_cluster is None:
            next_cluster = 0xffffffff
        value = struct.pack("<I", 0xffffffff & next_cluster)
    fat_container_file_fp.seek((fat_header_info['reserved_sec'] * fat_header_info['bytes_per_sec']) + offset)
    fat_container_file_fp.write(value)
    fat_container_file_fp.seek(((fat_header_info['reserved_sec'] + fat_header_info['sec_per_fat']) * fat_header_info[
        'bytes_per_sec']) + offset)
    fat_container_file_fp.write(value)
    """
    fat_container_file_fp.seek((fat_header_info['reserved_sec']*fat_header_info['bytes_per_sec'])+next_cluster)
    if fat_type.lower()=="fat16".lower():
        fat_container_file_fp.write(struct.pack("<h", 0xffff))
    if fat_type.lower()=="fat32".lower():
        fat_container_file_fp.write(struct.pack("<i", 0xffffffff))
    """


def get_cluster_chain(start_cluster_num):
    cluster_table = get_cluster_table()
    cluster_chain = [start_cluster_num]
    if fat_type.lower() == "fat16":
        valid_lower_bound = 0x0002
        valid_upper_bound = 0xFFEF
        # end_lower_bound = 0xFFF8
        # end_upper_bound = 0xFFFF
    if fat_type.lower() == "fat32":
        valid_lower_bound = 0x00000002
        valid_upper_bound = 0x0FFFFFEF
        # end_lower_bound = 0x0FFFFFF8
        # end_upper_bound = 0x0FFFFFFF
    while cluster_table[start_cluster_num] >= valid_lower_bound and cluster_table[
        start_cluster_num] <= valid_upper_bound:
        cluster_chain.append(cluster_table[start_cluster_num])
        start_cluster_num = cluster_table[start_cluster_num]
    return cluster_chain


def add_file_entry(real_file_name, file_name, dir_cluster=None, dir_table_start=None, dir_table_length=None):
    short_name = generate_short_name(file_name)
    check = short_name.replace('.', '')
    # Check to see if we're adding to the root directory of FAT16
    if fat_type.lower() == "fat16" and dir_cluster is None:
        dir_entry = find_directory_entry(check, None, dir_table_start, dir_table_length)
        if dir_entry is None:
            print "Root directory table is full."
            sys.exit(1)

    free_cluster = find_free_cluster()
    update_fat_cluster(free_cluster)
    print "Adding file", file_name

    dot_index = short_name.rfind(".")
    if dot_index != -1:
        file_part_name = short_name[:dot_index]
        ext_part_name = short_name[dot_index + 1:dot_index + 1 + 3]
    else:
        file_part_name = short_name
        ext_part_name = ''

    file_size = os.path.getsize(real_file_name)
    atime = time.localtime(os.path.getatime(real_file_name))
    ctime = time.localtime(os.path.getctime(real_file_name))
    mtime = time.localtime(os.path.getmtime(real_file_name))

    # Generate short name entry
    if fat_type.lower() == "fat16":
        dirent = fat_dirent.pack(file_part_name + ('\x20' * (8 - len(file_part_name))),
                                 ext_part_name + ('\x20' * (3 - len(ext_part_name))), 0x20, 0x18, 0x0,
                                 ((ctime[3] << 11) & 0xF800) | ((ctime[4] << 5) & 0x7E0) | ((ctime[5] / 2) & 0x1f),
                                 (((ctime[0] - 1980) << 9) & 0xFE00) | ((ctime[1] << 5) & 0x1E0) | (ctime[2] & 0x1F),
                                 (((atime[0] - 1980) << 9) & 0xFE00) | ((atime[1] << 5) & 0x1E0) | (atime[2] & 0x1F),
                                 0x0000,
                                 ((mtime[3] << 11) & 0xF800) | ((mtime[4] << 5) & 0x7E0) | ((mtime[5] / 2) & 0x1F),
                                 (((mtime[0] - 1980) << 9) & 0xFE00) | ((mtime[1] << 5) & 0x1E0) | (mtime[2] & 0x1F),
                                 free_cluster, file_size)
    if fat_type.lower() == "fat32":
        dirent = fat_dirent.pack(file_part_name + ('\x20' * (8 - len(file_part_name))),
                                 ext_part_name + ('\x20' * (3 - len(ext_part_name))), 0x20, 0x18, 0x0,
                                 ((ctime[3] << 11) & 0xF800) | ((ctime[4] << 5) & 0x7E0) | ((ctime[5] / 2) & 0x1f),
                                 (((ctime[0] - 1980) << 9) & 0xFE00) | ((ctime[1] << 5) & 0x1E0) | (ctime[2] & 0x1F),
                                 (((atime[0] - 1980) << 9) & 0xFE00) | ((atime[1] << 5) & 0x1E0) | (atime[2] & 0x1F),
                                 (free_cluster & 0xffff0000) >> 16,
                                 ((mtime[3] << 11) & 0xF800) | ((mtime[4] << 5) & 0x7E0) | ((mtime[5] / 2) & 0x1F),
                                 (((mtime[0] - 1980) << 9) & 0xFE00) | ((mtime[1] << 5) & 0x1E0) | (mtime[2] & 0x1F),
                                 free_cluster & 0x00000ffff, file_size)
    block_size = fat_header_info['sec_per_cluster'] * fat_header_info['bytes_per_sec']

    # Generate long name entry
    num_entries = round_up(len(file_name), 13)
    long_entries = GenerateLongEntry(file_name, short_name, num_entries)

    offset = None
    if fat_type.lower() == "fat16" and dir_cluster is None:
        fat_container_file_fp.seek(dir_table_start)
        dir_cluster_data = fat_container_file_fp.read(dir_table_length)
        for i in range(dir_table_start, dir_table_start + dir_table_length, DIR_ENT_LENGTH):
            dir_data = fat_dirent.unpack(dir_cluster_data[i - dir_table_start:i - dir_table_start + DIR_ENT_LENGTH])
            temp_name = dir_data[0].strip() + '.' + dir_data[1]
            # Make sure we aren't adding a duplicate name
            if short_name == temp_name:
                print "File with the same name already exists in directory."
                sys.exit(1)
            if dir_data[0].strip('\x00').strip() == "":
                if len(long_entries) > 0:
                    fat_container_file_fp.seek(i)
                    fat_container_file_fp.write(long_entries[-1])
                    del long_entries[-1]
                else:
                    offset = i
                    break
    else:
        cluster_chain = get_cluster_chain(dir_cluster)
        for each_cluster in cluster_chain:
            dir_cluster_data = read_cluster(each_cluster)
            for i in range(0, len(dir_cluster_data), DIR_ENT_LENGTH):
                dir_data = fat_dirent.unpack(dir_cluster_data[i:i + DIR_ENT_LENGTH])
                temp_name = dir_data[0].strip() + '.' + dir_data[1]
                # Make sure we aren't adding a duplicate name
                if short_name == temp_name:
                    print "File with the same name already exists in directory."
                    sys.exit(1)
                if dir_data[0].strip('\x00').strip() == "":
                    if len(long_entries) > 0:
                        fat_container_file_fp.seek(i + (get_cluster_sector(each_cluster)*fat_header_info['bytes_per_sec']))
                        fat_container_file_fp.write(long_entries[-1])
                        del long_entries[-1]
                    else:
                        offset = i + (get_cluster_sector(each_cluster)*fat_header_info['bytes_per_sec'])
                        break

    if offset is None:
        temp_next_cluster = find_free_cluster()
        update_fat_cluster(cluster_chain[-1], temp_next_cluster)
        update_fat_cluster(temp_next_cluster)
        write_cluster(temp_next_cluster, '')
        offset = (get_cluster_sector(temp_next_cluster) * fat_header_info['bytes_per_sec'])

    if not offset is None:
        # Finish writing the long entries
        for entry in reversed(long_entries):
            fat_container_file_fp.seek(offset)
            fat_container_file_fp.write(entry)
            offset = offset + 32
        fat_container_file_fp.seek(offset)
        fat_container_file_fp.write(dirent)

        try:
            fp = open(real_file_name, 'rb')
        except:
            PrintBigError("Can't open '%s' in binary read mode" % (real_file_name))

        while file_size > 0:
            write_cluster(free_cluster, fp.read(min(block_size, file_size)))
            file_size = file_size - min(block_size, file_size)
            if file_size > 0:
                new_free_cluster = find_free_cluster()
                update_fat_cluster(free_cluster, new_free_cluster)
                free_cluster = new_free_cluster
            update_fat_cluster(free_cluster)
        try:
            fp.close()
        except:
            PrintBigError("Can't close '%s'" % (real_file_name))
    print "Added ", file_name


def add_directory_entry(dir_name, short_name, dir_cluster=None, dir_table_start=None, dir_table_length=None):
    global chain_start
    print "Adding directory ", dir_name

    dot_index = short_name.rfind(".")
    if dot_index != -1:
        file_part_name = short_name[:dot_index]
        ext_part_name = short_name[dot_index + 1:dot_index + 1 + 3]
    else:
        file_part_name = short_name
        ext_part_name = ''

    # Generate long name entry
    num_entries = round_up(len(dir_name), 13)
    long_entries = GenerateLongEntry(dir_name, short_name, num_entries)

    offset = None
    # If we're adding to a root directory of fat16
    if fat_type.lower() == "fat16" and dir_cluster is None:
        fat_container_file_fp.seek(dir_table_start)
        dir_cluster_data = fat_container_file_fp.read(dir_table_length)
        for i in range(dir_table_start, dir_table_start + dir_table_length, DIR_ENT_LENGTH):
            dir_data = fat_dirent.unpack(dir_cluster_data[i - dir_table_start:i - dir_table_start + DIR_ENT_LENGTH])
            if dir_data[0].strip('\x00').strip() == "":
                if len(long_entries) > 0:
                    fat_container_file_fp.seek(i)
                    fat_container_file_fp.write(long_entries[-1])
                    del long_entries[-1]
                else:
                    offset = i
                    break
    else:
        dir_cluster_data = read_cluster(dir_cluster)
        for i in range(0, len(dir_cluster_data), DIR_ENT_LENGTH):
            dir_data = fat_dirent.unpack(dir_cluster_data[i:i + DIR_ENT_LENGTH])
            if dir_data[0].strip('\x00').strip() == "":
                if len(long_entries) > 0:
                    fat_container_file_fp.seek(
                        i + (get_cluster_sector(dir_cluster) * fat_header_info['bytes_per_sec']))
                    fat_container_file_fp.write(long_entries[-1])
                    del long_entries[-1]
                else:
                    offset = i + (get_cluster_sector(dir_cluster) * fat_header_info['bytes_per_sec'])
                    break

    if offset is None:
        cluster_chain = get_cluster_chain(current_cluster)
        temp_next_cluster = find_free_cluster()
        update_fat_cluster(cluster_chain[-1], temp_next_cluster)
        update_fat_cluster(temp_next_cluster)
        write_cluster(temp_next_cluster, '')
        offset = (get_cluster_sector(temp_next_cluster) * fat_header_info['bytes_per_sec'])

    if not offset is None:
        # Finish writing the long entries
        for entry in reversed(long_entries):
            fat_container_file_fp.seek(offset)
            fat_container_file_fp.write(entry)
            offset = offset + 32
        free_cluster = find_free_cluster()
        dirent = fat_dirent.pack(file_part_name + ('\x20' * (8 - len(file_part_name))),
                                 ext_part_name + ('\x20' * (3 - len(ext_part_name))), 0x10, 0x08, 0x80, 0x70EA, 0x3EE7,
                                 0x3EE7, 0x0000, 0x8BC8, 0x3EE5, free_cluster, 0)
        fat_container_file_fp.seek(offset)
        fat_container_file_fp.write(dirent)

        # Always set .. to the start of the cluster chain
        write_cluster(free_cluster,
                      fat_dirent.pack('.' + ('\x20' * (8 - len('.'))), 3 * '\x20', 0x10, 0, 0x93, 0x70EA, 0x3EE7,
                                      0x3EE7, 0x0000, 0x70EB, 0x3EE7, free_cluster, 0) + \
                      fat_dirent.pack('..' + ('\x20' * (8 - len('..'))), 3 * '\x20', 0x10, 0, 0x93, 0x70EA, 0x3EE7, 0x3EE7, 0x0000,
                                      0x70EB, 0x3EE7, chain_start, 0))
        update_fat_cluster(free_cluster)
        chain_start = free_cluster
    print "Added ", dir_name
    return free_cluster


def find_directory_entry(dir_name, dir_cluster=None, dir_table_start=None, dir_table_length=None):
    global current_cluster
    global chain_start
    dir_name = dir_name.replace(".", '')
    if fat_type.lower() == "fat16" and dir_cluster is None:
        fat_container_file_fp.seek(dir_table_start)
        dir_cluster_data = fat_container_file_fp.read(dir_table_length)
        return find_directory_entry_helper(dir_name, dir_cluster_data)
    else:
        cluster_chain = get_cluster_chain(dir_cluster)
        if dir_cluster != 2:
            chain_start = dir_cluster
        else:
            chain_start = 0
        for each_cluster in cluster_chain:
            result = find_directory_entry_helper(dir_name, read_cluster(each_cluster))
            if not result is None:
                if result is False:
                    current_cluster = each_cluster
                return result
        current_cluster = each_cluster
    return None


def find_directory_entry_helper(dir_name, dir_cluster_data):
    """
    True: directory exists, current_cluster holds the cluster number of the directory
    False: directory does not exist, but there's space in the current searched cluster for it
    None: directory does not exist, no space found for it in the current searched cluster
    """
    global current_cluster
    for i in range(0, len(dir_cluster_data), DIR_ENT_LENGTH):
        dir_data = fat_dirent.unpack(dir_cluster_data[i:i + DIR_ENT_LENGTH])
        if (dir_data[0].strip('\x00').strip().upper() + dir_data[1].strip('\x00').strip().upper()) == dir_name:
            current_cluster = dir_data[11]
            return True
        if dir_data[0].strip('\x00').strip() == "":
            return False
    return None


def add_entry(fat_file_name, real_file_name):
    global current_cluster
    global chain_start
    fat_entry_list = []
    # Create a list out of entry name
    while fat_file_name != '':
        temp = os.path.split(fat_file_name)
        fat_entry_list.insert(0, temp[1].strip().strip("."))
        fat_file_name = temp[0]
    #fat_entry_list = map(clean_short_dos_name, fat_entry_list)
    # Always starts at the root directory
    root_entry = True
    chain_start = 0
    if fat_type.lower() == "fat16":
        # zero-indexed start address
        reserved_sec = fat_header_info['reserved_sec']
        num_fat_copies = fat_header_info['num_fat_copies']
        sec_per_fat = fat_header_info['sec_per_fat']
        bytes_per_sec = fat_header_info['bytes_per_sec']
        dir_table_start = (reserved_sec + (num_fat_copies * sec_per_fat)) * bytes_per_sec
        # length in bytes (of this chunk of directory info)
        dir_table_length = (fat_header_info['max_root_dir_ent'] * 32)
    else:
        current_cluster = fat_header_info['root_dir_start_cluster']
    # Loop through entry list
    while fat_entry_list:
        if len(fat_entry_list[0]) > 255:
            print "Entry is longer than 255 characters."
            sys.exit(1)
        # import pdb; pdb.set_trace()
        # If it's a single file, this will only contain 1 item
        if 1 == len(fat_entry_list) and os.path.isfile(real_file_name):
            # add file
            if root_entry and fat_type.lower() == "fat16":
                add_file_entry(real_file_name, fat_entry_list[0], None, dir_table_start, dir_table_length)
            else:
                add_file_entry(real_file_name, fat_entry_list[0], current_cluster)
        # If it's a directory, it will have multiple items
        else:
            short_name = generate_short_name(fat_entry_list[0])
            # add dir
            if root_entry and fat_type.lower() == "fat16":
                # Look for directory entry, don't create it if it exists
                dir_entry = find_directory_entry(short_name, None, dir_table_start, dir_table_length)
                # Stop the process if the root directory is full in fat16
                if dir_entry is None:
                    print "Root directory table is full."
                    sys.exit(1)
                elif dir_entry == False:
                    current_cluster = add_directory_entry(fat_entry_list[0],short_name, None, dir_table_start, dir_table_length)
            else:
                dir_entry = find_directory_entry(short_name, current_cluster)
                # Could not find room for our directory in the current cluster, look for a different one
                if dir_entry is None:
                    cluster_chain = get_cluster_chain(current_cluster)
                    free_cluster = find_free_cluster()
                    update_fat_cluster(cluster_chain[-1], free_cluster)
                    update_fat_cluster(free_cluster)
                    current_cluster = add_directory_entry(fat_entry_list[0], short_name, free_cluster)
                elif dir_entry == False:
                    current_cluster = add_directory_entry(fat_entry_list[0], short_name, current_cluster)

        del fat_entry_list[0]
        if root_entry:
            root_entry = False


def usage():
    print """
Usage: python fatadd.py [OPTION...]
Add files to a FAT16 or FAT32 container.

Examples:
  python fatadd.py --name=FAT16_256MB.bin --from=C:\TEMP\                  # Add files and directories recursively from C:\TEMP into the
                                                                           # root directory of the container FAT16_256MB.bin
  python fatadd.py --name=FAT16_256MB.bin --from=C:\TEMP\ --dir=images     # Add files and directories recursively from C:\TEMP into the 
                                                                           # images directory of the container FAT16_256MB.bin
  python fatadd.py --name=FAT16_256MB.bin --from=C:\TEMP\myfile.txt        # Add C:\TEMP\myfile.txt into the
                                                                           # root directory of the container FAT16_256MB.bin
  Any non-existent directories along the way will be created.

 Options:

  -d, --dir                  Directory to begin with (root, if omitted)
  -n, --name                 Container name
  -f, --from                 File/directory to add into the container
  -c, --sectorsize           Sector size to use
  -?, --help                 Display this help message
  -v, --verbose              Verbose messages

    """


def main():
    print "fatadd version 1.11"
    global fat_container_file, fat_type, fat_container_file_fp, SECTOR_SIZE

    dir_name = None
    entry_name = None

    try:
        opts, args = getopt.getopt(sys.argv[1:], "d:n:f:c:?v",
                                   ["dir=", "name=", "from=", "sectorsize=", "help", "verbose"])
    except getopt.GetoptError, err:
        # print help information and exit:
        print str(err)  # will print something like "option -a not recognized"
        usage()
        sys.exit(1)

    for o, a in opts:
        if o in ("-d", "--dir"):
            dir_name = a.decode("utf-8")
        elif o in ("-n", "--name"):
            fat_container_file = a
        elif o in ("-f", "--from"):
            entry_name = a.decode("utf-8")
            # entry_name = a
        elif o in ("-c", "--sectorsize"):
            SECTOR_SIZE = int(a)
        elif o in ("-?", "--help"):
            usage()
            sys.exit(0)
        elif o in ("-v", "--verbose"):
            pass
        else:
            assert False, "unhandled option"

    if fat_container_file is None:
        usage()
        print "ERROR: No FAT container provided."
        sys.exit(1)
    if entry_name is None:
        usage()
        print "ERROR: No entries were provided to add to container."
        sys.exit(1)
    if dir_name is None:
        dir_name = ''

    # Get all the header information from FAT image
    read_header()

    try:
        fat_container_file_fp = open(fat_container_file, 'r+b')
    except:
        PrintBigError("Can't open '%s' in binary read/write mode" % fat_container_file)

    if fat_type.lower() == "fat32":
        cluster_chain = get_cluster_chain(fat_header_info['root_dir_start_cluster'])
        # If nothing in the FAT table has been linked yet, update
        if len(cluster_chain) == 1:
            update_fat_cluster(fat_header_info['root_dir_start_cluster'])
        # If we're adding a file
    if os.path.isfile(entry_name):
        add_entry(os.path.join(dir_name, os.path.basename(entry_name)), entry_name)
    else:
        # If we're adding a sub directory
        entry_name = os.path.normpath(entry_name)
        # Walk through and add every file
        for root, dirs, files in os.walk(entry_name, topdown=False):
            for f in files:
                real_file_name = os.path.join(root, f)
                fat_file_name = real_file_name[real_file_name.index(entry_name) + len(entry_name):]
                fat_file_name = fat_file_name.strip("\\").strip("/")
                add_entry(os.path.join(dir_name, fat_file_name), real_file_name)
            for d in dirs:
                real_file_name = os.path.join(root, d)
                fat_file_name = real_file_name[real_file_name.index(entry_name) + len(entry_name):]
                fat_file_name = fat_file_name.strip("\\").strip("/")
                add_entry(os.path.join(dir_name, fat_file_name), real_file_name)

    fat_container_file_fp.close()


if __name__ == "__main__":
    main()