firmex/python/main.py

import argparse
import os
import io

import matcher

signatures = [
    matcher.Zip,
    matcher.Ambarella,
    matcher.SquashFS,
    matcher.RomFS,
    matcher.CromFS,
    matcher.FlattenDeviceTree,
    matcher.ELF,
]

def detect(args):
    print('detecting', args.file)
    matches = []

    # recursive?
    for matcher in signatures:
        m = matcher(args.file)
        if m.is_valid():
            matches += [m]

    for filetype in matches:
        print("detected", filetype.name)
        for m in filetype.matches:
            print(">", m)
            filetype.view(m)

    return matches


def extract(args):
    pass

def main():
    parser = argparse.ArgumentParser(description='Program for detecting or extracting data.')

    subparsers = parser.add_subparsers(dest='command')

    # Subparser for the 'detect' command
    detect_parser = subparsers.add_parser('detect', help='Detect data in a file.')
    detect_parser.add_argument('file', help='Input file')
    detect_parser.add_argument('--isa', action='store_true', help='Perform ISA detection')

    # Subparser for the 'extract' command
    extract_parser = subparsers.add_parser('extract', help='Extract data from a file.')
    extract_parser.add_argument('file', help='Input file')
    extract_parser.add_argument('--dry', action='store_true', help='Perform a dry run without extracting')

    args = parser.parse_args()

    if args.command is None:
        print('no command given')
        exit(1)

    if args.file is None:
        print('require file')
        exit(1)

    if not os.path.exists(args.file) or os.path.isdir(args.file):
        print('please provide an existing file')
        exit(1)


    if args.command == 'detect':
        detect(args)
        # if args.isa:
        #     # Perform ISA detection on the file
        #     print('Performing ISA detection on:', args.file)
        # else:
        #     parser.print_help()
    elif args.command == 'extract':
        extract(args)
        # if args.dry:
        #     # Perform a dry run without extracting
        #     print('Dry run extraction from:', args.file)
        # else:
        #     # Extract data from the file
        #     print('Extracting data from:', args.file)
    else:
        parser.print_help()

if __name__ == '__main__':
    main()