Corrupt FAT directory entries

corrupt.py

#!/usr/bin/env python3

# prepare image:
#  dd if=/dev/zero of=img bs=1M count=1
#  mkfs.vfat img
#  echo 'foo=bar' > uboot.env
#  mcopy -i img uboot.env ::uboot.env
#  mcopy -i img uboot.env ::'uboot canary long name.env'
#  ./corrupt.py img --corrupt 'uboot canary long name.env' --short=uboot.env --long=uboot.env --update-checksum
#  /snap/core/current/sbin/fsck.vfat -av img
#  mdir -i img



import struct
import argparse


def parse_arguments():
    parser = argparse.ArgumentParser(description='corrupt file in FAT')
    parser.add_argument('path', help='path')
    parser.add_argument('-c', '--corrupt',
                        help='corrupt directory entries with this long name',
                        type=str)
    parser.add_argument('--long',
                        help='change long name to this string',
                        default='',
                        type=str)
    parser.add_argument('--short',
                        help='change short name to this string',
                        default='',
                        type=str)
    parser.add_argument('--update-checksum',
                        help='update LFN checksum', action='store_true',
                        default=False)
    return parser.parse_args()


def c2f(count):
    if count == 1:
        bformat = '<B'  # 1 byte
    elif count == 2:
        bformat = '<H'  # halfword
    elif count == 4:
        bformat = '<i'  # word
    else:
        raise ValueError('unsupported byte count {}'.format(count))
    return bformat


def readbytes(inf, offs, count):
    inf.seek(offs)
    return  struct.unpack(c2f(count), inf.read(count))[0]


def readstring(inf, offs, count, encoding='ascii'):
    inf.seek(offs)
    raw = inf.read(count)
    # print('raw string', raw)
    try:
        return struct.unpack('{}s'.format(count), raw)[0].decode(encoding)
    except UnicodeDecodeError:
        return '<decoding failed, raw: {}>'.format(raw)


def writestring(inf, offs, count, what, encoding='ascii'):
    inf.seek(offs)
    if len(what) > count:
        what = what[0:count]
    else:
        what += ' ' * (count - len(what))
    print('writing "{}" len: {}'.format(what, len(what)))
    inf.write(what.encode(encoding))


def writebytes(inf, offs, count, what):
    inf.seek(offs)
    if type(what) is bytes:
        bwhat = what[0:count]
        try:
            print('writing:', bwhat.decode('utf-16'), '--', what.decode('utf-16'))
        except UnicodeError:
            pass
    else:
        bwhat = struct.pack(c2f(count), what)
    inf.write(bwhat)


def lfncksum(name, ext):
    if len(name) > 8 or len(ext) > 3:
        raise ValueError('name or length of incorrect size')
    ckname = name + ' '*(8 - len(name)) + ext + ' ' * (3 - len(ext))
    # print('ckname "{}"'.format(ckname))
    psum = 0
    for lttr in ckname:
        psum = (((psum & 1) << 7) + (psum >> 1) + ord(lttr)) & 0xff
        # print('psums =', psum, 'c =', lttr, 'ord(lttr) =', ord(lttr))
    # print('checksum for', name, ext, 'is', psum)
    return psum


def main(path, corrupt, newshort, newlong, updatecksum):
    with open(path, 'r+b') as inf:
        # bytes per logical sector
        bpls = readbytes(inf, 0xb, 2)
        # sectors per cluster
        spc = readbytes(inf, 0xd, 1)
        # reserved sectors
        reserved = readbytes(inf, 0xe, 2)
        # # of FATs
        fats = readbytes(inf, 0x10, 1)
        # DOS4.x support
        dos4 = readbytes(inf, 0x26, 1) == 0x29
        # FAT32 indicator
        fat32 = readbytes(inf, 0x42, 1) in [0x28, 0x29]
        if fat32:
            # sectors per FAT
            spf = readbytes(inf, 0x24, 4)
            # filesystem type - FAT32
            fstype = readstring(inf, 0x52, 8).strip()
            # label
            label = readstring(inf, 0x47, 11).strip()
            maxdirentries = 0
        else:
            # sectors per FAT
            spf = readbytes(inf, 0x16, 2)
            # filesystem type - FAT12|FAT16
            fstype = readstring(inf, 0x36, 8).strip()
            # label
            label = readstring(inf, 0x2b, 11).strip()
            # maximum # of directory entries
            maxdirentries = readbytes(inf, 0x11, 2)

        print('bytes per sector', bpls)
        print('sectors per cluster', spc)
        print('# of fats', fats)
        print('size per fat', spf)
        print('fat 32', fat32)
        print('fstype', fstype)
        print('dos4', dos4)
        print('label', label)
        print('maxdirentries', maxdirentries)

        for i in range(1, fats + 1):
            start = (reserved + (i - 1) * spf) * bpls
            print('fat #{0} start 0x{1:x}'.format(i, start))

        rootdirstart = (reserved + fats * spf) * bpls
        print('rootdir start 0x{0:x}'.format(rootdirstart))

        entry = 0
        longname = ''

        while maxdirentries != 0 and entry < maxdirentries:
            entrystart = rootdirstart + entry * 32
            # short name
            name = readstring(inf, entrystart, 8)
            # is it empty?
            if name[0] in ['\x00', '\xe5']:
                entry += 1
                continue
            # read only, hidden, system, volume bits set if entry is a VFAT
            # long name
            islongname = readbytes(inf, entrystart + 0xb, 1) == 0xf
            if islongname:
                # long name entry, they precede the corresponding short name
                # entry

                # seqn = readbytes(inf, entrystart, 1)
                # pick up respective pieces of long name
                name = readstring(inf, entrystart + 0x1, 10, 'utf-16') + \
                       readstring(inf, entrystart + 0xe, 12, 'utf-16') + \
                       readstring(inf, entrystart + 0x1c, 4, 'utf-16')
                longname = name + longname

            else:
                # got shortname entry, also it marks a finish of any preceding
                # long name entries

                # extension
                ext = readstring(inf, entrystart + 0x8, 3)
                shortname = '{}.{}'.format(name.strip(), ext)

                print('- {}.{}\tlong: {}'.format(name.strip(),
                                                 ext,
                                                 longname if longname else shortname + ' (from short)'))
                if corrupt and (corrupt == shortname or corrupt == longname):
                    preventrystart = entrystart - 32

                    if newshort:
                        print('corrupt match, set shortname to', newshort.upper())
                        shortbase, shortext = newshort.upper().split('.')
                        writestring(inf, entrystart, 8, shortbase)
                        writestring(inf, entrystart + 0x8, 3, shortext)
                        name, ext = shortbase, shortext

                    # long name enrty precedes short name one
                    isprevlongname = readbytes(inf, preventrystart + 0xb, 1) == 0xf

                    if newlong:
                        assert isprevlongname, 'previous entry at offset 0x{0:x} not a LFN'.format(preventrystart)

                        print('corrupt match, set longname to', newlong)
                        lfn = newlong.encode('utf-16')
                        if len(lfn) < 26:
                            lfn += b'\x00\x00'
                            lfn += b'\xff\xff' * ((26 - len(lfn)) // 2)
                        writebytes(inf, preventrystart + 0x1, 10, lfn)
                        writebytes(inf, preventrystart + 0xe, 12, lfn[10:])
                        writebytes(inf, preventrystart + 0x1c, 4, lfn[22:])

                    if longname and updatecksum:
                        assert isprevlongname, 'previous entry at offset 0x{0:x} not a LFN'.format(preventrystart)
                        cksum = lfncksum(name, ext)
                        print('update LFN cksum to', cksum)
                        writebytes(inf, preventrystart + 0xd, 1, cksum)

                longname = ''

            entry += 1


if __name__ == '__main__':
    opts = parse_arguments()
    print('corrupt', opts.corrupt)
    if opts.corrupt and (not opts.short and not opts.long):
        print('use --short or --long to pass the desired long and short names')
        raise SystemExit(1)
    if len(opts.long) > 13:
        print('LFN longer than 13 characters is not supported')
        raise SystemExit(1)
    main(opts.path, opts.corrupt, opts.short, opts.long, opts.update_checksum)