main.cpp

#include "finders/finder.hpp"
#include "finders/zero_finder.hpp"
#include "finders/little2_hash_finder.hpp"
#include "finders/md5_hash_finder.hpp"
#include "finders/ext4_finder.hpp"
#include "finders/veeam_page_finder.hpp"
#include "finders/veeam_slot_finder.hpp"
#include "finders/zfs_finder.hpp"

#include "readers/reader.hpp"

#include "file_or_dev_size.hpp"

#include <fcntl.h>
#include <getopt.h>
#include <limits.h>
#include <sys/ioctl.h>
#include <sys/signal.h>
#include <unistd.h>

enum Mode {
    FIND_PATTERN,
    FIND_ZEROES,
    FIND_L2_HASH,
    FIND_MD5_HASH,
    FIND_EXT4,
    FIND_VEEAM_PAGES,
    FIND_VEEAM_SLOTS,
    FIND_ZFS_BLKPTR,
};

enum PatternType {
    PT_HEX_BYTES,
    PT_HEX_INT,
    PT_ASCII_STRING,
    PT_UNICODE_STRING,
    PT_ANY_STRING
};

bool g_stop = false;

void signalHandler(int signum) {
    // extra linefeed to keep the last progress line
    printf("\n[^] Interrupt signal (%d) received. Exiting gracefully...\n", signum);
    g_stop = true;
    Finder::stop();
}

std::string ascii2unicode(const std::string& str) {
    std::string unicode_str;
    for (char c : str) {
        unicode_str.push_back(c);
        unicode_str.push_back(0); // add null byte for 16-bit encoding
    }
    return unicode_str;
}

void print_usage(int retcode) {
    printf("Usage: binsearch [options] <image_or_device_fname> <pattern1> .. [patterN]\n");
    printf("  -t, --pattern-type TYPE        [bin|int|ascii|unicode|text] (default: bin)\n");
    printf("                                 bin:     binary data (hex)\n");
    printf("                                 int:     reverses byte order\n");
    printf("                                 ascii:   search for arbitrary ASCII text\n");
    printf("                                 unicode: convert argument to 16-bit encoding, i.e. UCS-2\n");
    printf("                                 text:    try all supported encodings\n");
    puts("");
    printf("  -h, --help                     show this help message and exit\n");
    printf("      --start OFFSET             start search from specified offset (default: 0)\n");
    printf("      --block-size BLOCK_SIZE    calculate block-based stuff using this block size (default: None)\n");
    printf("      --step SIZE                advance by SIZE bytes (default: 1)\n");
    printf("      --shift OFFSET             look for the pattern at specified OFFSET in the block (default: 0)\n");
    printf("  -H, --hexdump-width WIDTH      number of data bytes to show when pattern is found (default: %d)\n", DEFAULT_HEXDUMP_WIDTH);
    printf("\n");
    printf("  -w  --write-bin-offsets FNAME  write offsets of found data as uint64 to specified file\n");
    printf("  -a  --append                   append to file (default: overwrite)\n");
    printf("\n");
    printf("  -q, --quiet                    decrease verbosity\n");
    printf("  -v, --verbose                  increase verbosity\n");
    printf("\n");
    printf("  -z  --find-zeroes              find continuous zero regions instead of searching for a pattern\n");
    printf("      --find-l2hash              find the block with specified little2 hash\n");
    printf("      --find-md5                 find the block with specified md5 hash\n");
    printf("      --find-ext4                find ext4 extents\n");
    printf("      --find-veeam-pages         find veeam 4kb VBlob pages prefixed with MD5 hash\n");
    printf("      --find-veeam-slots         find veeam VBK CSlotHdr prefixed with vcrc32\n");
    printf("\n");
    printf("      --find-zfs-blkptr          find ZFS blkptr_t\n");
    printf("      --zfs-blkptr-level LVL     find only blkptrs of this level (default: any level)\n");
    printf("      --zfs-blkptr-max-psize X   find only blkptrs smaller than X (default: any size)\n");
    exit(retcode);
}

int main(int argc, char* argv[]) {
    static struct option long_options[] = {
        {"help",              no_argument, 0, 'h'},
        {"verbose",           no_argument, 0, 'v'},
        {"quiet",             no_argument, 0, 'q'},

        {"write-bin-offsets", required_argument, 0, 'w'},
        {"append",            no_argument, 0, 'a'},

        {"pattern-type",      required_argument, 0, 't'},
        {"start",             required_argument, 0, 0},
        {"step",              required_argument, 0, 0},
        {"block-size",        required_argument, 0, 0},
        {"shift",             required_argument, 0, 0},
        {"hexdump-width",     required_argument, 0, 'H'},

        {"find-zeroes",       no_argument, 0, 'z'},
        {"find-l2hash",       no_argument, 0, 0},
        {"find-md5",          no_argument, 0, 0},
        {"find-ext4",         no_argument, 0, 0},
        {"find-veeam-pages",  no_argument, 0, 0},
        {"find-veeam-slots",  no_argument, 0, 0},

        {"find-zfs-blkptr",   optional_argument, 0, 0},
        {"zfs-blkptr-level",  required_argument, 0, 0},
        {"zfs-blkptr-max-psize", required_argument, 0, 0},

        {0, 0, 0, 0}
    };
    int opt;
    size_t start = 0, step = 1, shift = 0, block_size = 0;
    Mode mode = FIND_PATTERN;
    PatternType pattern_type = PT_HEX_BYTES;
    int verbosity = 0;
    const char* offsets_fname = nullptr;
    std::string offsets_mode = "wb";
    int option_index = 0;
    int hexdump_width = DEFAULT_HEXDUMP_WIDTH;
    std::unique_ptr<Finder> finder;

    while ((opt = getopt_long(argc, argv, "ahH:qt:vw:z", long_options, &option_index)) != -1) {
        switch (opt) {
            case 'a': offsets_mode = "ab"; break;
            case 'h': print_usage(EXIT_SUCCESS); break;
            case 'H': {
                char* end;
                long val = strtol(optarg, &end, 10);
                if (*end != '\0' || val <= 0 || val > 65536) {
                    fprintf(stderr, "Error: invalid hexdump width: %s\n", optarg);
                    exit(1);
                }
                hexdump_width = (int)val;
                break;
            }
            case 'v': verbosity++; break;
            case 'w': offsets_fname = optarg; break;
            case 'q': verbosity--; break;
            case 't': 
                  switch( *optarg ) {
                      case 'i':
                          pattern_type = PT_HEX_INT;
                          break;
                      case 'a':
                          pattern_type = PT_ASCII_STRING;
                          break;
                      case 'u':
                          pattern_type = PT_UNICODE_STRING;
                          break;
                      case 't':
                          pattern_type = PT_ANY_STRING;
                          break;
                      case 'b':
                          pattern_type = PT_HEX_BYTES;
                          break;
                      default:
                          print_usage(EXIT_FAILURE);
                          break;
                  }
                break;
            case 'z':
                mode = FIND_ZEROES;
                if( step == 1 ){
                    step = ZERO_CHUNK_SIZE;
                }
                break;
            case 0:
                if (strcmp(long_options[option_index].name, "start") == 0)
                    start = human2bytes(optarg);
                if (strcmp(long_options[option_index].name, "step") == 0)
                    step = human2bytes(optarg);
                if (strcmp(long_options[option_index].name, "block-size") == 0)
                    block_size = human2bytes(optarg);
                if (strcmp(long_options[option_index].name, "shift") == 0)
                    shift = human2bytes(optarg);
                if (strcmp(long_options[option_index].name, "find-l2hash") == 0)
                    mode = FIND_L2_HASH;
                if (strcmp(long_options[option_index].name, "find-md5") == 0)
                    mode = FIND_MD5_HASH;
                if (strcmp(long_options[option_index].name, "find-ext4") == 0)
                    mode = FIND_EXT4;
                if (strcmp(long_options[option_index].name, "find-veeam-pages") == 0){
                    mode = FIND_VEEAM_PAGES;
                    if( step == 1 ){
                        step = VEEAM_PAGE_SIZE;
                    }
                }
                if (strcmp(long_options[option_index].name, "find-veeam-slots") == 0){
                    mode = FIND_VEEAM_SLOTS;
                    if( step == 1 ){
                        step = VEEAM_PAGE_SIZE;
                    }
                }
                if (strcmp(long_options[option_index].name, "find-zfs-blkptr") == 0){
                    mode = FIND_ZFS_BLKPTR;
                    finder.reset(new ZFSBlkPtrFinder());
                }
                if (strcmp(long_options[option_index].name, "zfs-blkptr-level") == 0){
                    if( !finder ){
                        printf("Error: --%s must be after --find-zfs-blkptr\n", long_options[option_index].name);
                        exit(1);
                    }
                    int level = atoi(optarg);
                    ((ZFSBlkPtrFinder*)finder.get())->set_level(level);
                }
                if (strcmp(long_options[option_index].name, "zfs-blkptr-max-psize") == 0){
                    if( !finder ){
                        printf("Error: --%s must be after --find-zfs-blkptr\n", long_options[option_index].name);
                        exit(1);
                    }
                    uint64_t max_psize = human2bytes(optarg);
                    ((ZFSBlkPtrFinder*)finder.get())->set_max_psize(max_psize);
                }
                break;
            default:
                print_usage(EXIT_FAILURE);
        }
    }

    // always need at least one positional argument
    if( argc - optind < 1 ){
        print_usage(EXIT_FAILURE);
    }

    std::vector<std::string> patterns;
    if( mode != FIND_ZEROES && mode != FIND_EXT4 && mode != FIND_VEEAM_PAGES && mode != FIND_VEEAM_SLOTS && mode != FIND_ZFS_BLKPTR ){
        if( argc - optind < 2 ){
            print_usage(EXIT_FAILURE);
        }

        for(int i=optind+1; i<argc; i++){
            patterns.push_back(argv[i]);
        }

        switch( pattern_type ){
            case PT_HEX_BYTES:
                for(int i=0; i<patterns.size(); i++){
                    patterns[i] = Finder::hex2bin(patterns[i]);
                }
                break;
            case PT_HEX_INT:
                for(int i=0; i<patterns.size(); i++){
                    std::string bin_pattern = Finder::hex2bin(patterns[i]);
                    std::reverse(bin_pattern.begin(), bin_pattern.end());
                    patterns[i] = bin_pattern;
                }
                break;
            case PT_ASCII_STRING:
                // keep as-is
                break;
            case PT_UNICODE_STRING:
                for(int i=0; i<patterns.size(); i++){
                    patterns[i] = ascii2unicode(patterns[i]);
                }
                break;
            case PT_ANY_STRING:
                {
                    size_t npatterns = patterns.size(); // size before conversion
                    for(int i=0; i<npatterns; i++){
                        patterns.push_back(ascii2unicode(patterns[i]));
                    }
                }
                break;
            default:
                fprintf(stderr, "Error: invalid pattern type\n");
                return 1;
                break;
        }

        for(int i=0; i<patterns.size(); i++){
            if( patterns[i].size() == 0 ){
                fprintf(stderr, "Error: empty pattern #%d\n", i);
                return 1;
            }
        }
    }

    const char *fname = argv[optind];
    int fd = open(fname, O_RDONLY);
    if (fd == -1) {
        fprintf(stderr, "Error opening %s: %s\n", fname, strerror(errno));
        return 1;
    }

    uint64_t size = getFileOrDeviceSize(fd);
    if (size == 0) {
        fprintf(stderr, "[?] getFileOrDeviceSize(%d) returned 0\n", fd);
        close(fd);
        return 1;
    }

    switch( mode ){
        case FIND_PATTERN:
            finder.reset(Finder::create_bin(patterns, step, shift));
            break;
        case FIND_ZEROES:
            finder.reset(new ZeroFinder(step));
            break;
        case FIND_L2_HASH:
            finder.reset(new Little2HashFinder(patterns));
            break;
        case FIND_MD5_HASH:
            finder.reset(new MD5HashFinder(patterns));
            break;
        case FIND_EXT4:
            finder.reset(new Ext4Finder());
            break;
        case FIND_VEEAM_PAGES:
            finder.reset(new VeeamPageFinder());
            break;
        case FIND_VEEAM_SLOTS:
            finder.reset(new VeeamSlotFinder());
            break;
        case FIND_ZFS_BLKPTR:
            break;
        default:
            throw std::invalid_argument("invalid mode");
    }

    FILE* offsets_file = nullptr;
    if( offsets_fname ){
        offsets_file = fopen(offsets_fname, offsets_mode.c_str());
        if( !offsets_file ){
            fprintf(stderr, "Error opening %s: %s\n", offsets_fname, strerror(errno));
            return 1;
        }
        finder->set_offsets_file(offsets_file);
    }

    finder->set_start_pos(start);
    finder->set_total_size(size);
    finder->set_step(step);
    finder->set_hexdump_width(hexdump_width);

    auto reader = Reader::create(fname, start, getDeviceAlign(fd));
    reader->set_overlap(finder->get_overlap());

    printf("[*] processing %s: reader=%s, finder=%s, start=%s, size=%s, step=%s, patterns=%zu",
        fname,
        reader->to_string().c_str(),
        finder->name(),
        bytes2human(start).c_str(),
        bytes2human(size).c_str(),
        bytes2human(step).c_str(),
        (size_t)patterns.size()
    );

    if( shift != 0 ){
        printf(", shift=%s", bytes2human(shift).c_str());
        finder->set_shift(shift);
    }
    if( block_size != 0 ){
        printf(", block_size=%s", bytes2human(block_size).c_str());
        finder->set_block_size(block_size);
    }
    if( offsets_file ){
        printf(", offsets_file=\"%s\", mode=\"%s\"", offsets_fname, offsets_mode.c_str());
    }
    puts("");
    fflush(stdout);

    finder->set_verbosity(verbosity);

    signal(SIGINT, signalHandler);

    Chunk chunk;
    finder->start();
    while (!g_stop && reader->next_chunk(chunk)){
        finder->update_progress(chunk.pos);
        finder->process_chunk(chunk);
    }
    finder->finish();

    if( offsets_file ){
        fclose(offsets_file);
        offsets_file = nullptr;
    }

    return 0;
}