# Invader
# https://github.com/cgio/invader


def fs(path, pattern, start_offset=0, chunk_size=-1, chunk_limit=-1,
       find_all=True):
    """
    Yields offset of found byte pattern within a file.

    Supports wildcard bytes, starting offset, reading in chunks, and read
    limits.

    Args:
        path (str): The file path.
        pattern (str): The sequence of bytes, e.g. 'FF??E82623D7'.
            '??' represents a single byte wildcard.
            Spaces between bytes are supported.
        start_offset (int): The offset to start searching from.
        chunk_size (int): The read length per chunk (-1 == entire file).
        chunk_limit (int): The max # of chunks to read (-1 == all chunks).
        find_all (bool): True == find all instances. False == only first
            instance.

    Returns:
        int: On success, yield found offset.
        bool: On error, yield False.

    Example:
        import invader

        for found_offset in invader.fs(
            r'C:\target.exe',
            '?? 01 55 ?? ?? 4B 20 1E 1D ?? 15',
            start_offset=0x1000,
            chunk_size=1024,
            chunk_limit=10,
            find_all=False
        ):
            if found_offset is not False:
                print(hex(found_offset))
    """

    # Only -1 or > 0 is allowed
    if chunk_size == 0 or chunk_limit == 0:
        yield False

    pattern = pattern.replace(' ', '')

    # If no path, invalid pattern, or pattern is all wildcards
    if len(path) == 0 or len(pattern) < 2 or len(
            pattern) % 2 != 0 or pattern.count('?') == len(pattern):
        yield False

    # Correct invalid values
    if start_offset < 0:
        start_offset = 0

    # If chunk_size == entire file, chunk_limit becomes irrelevant
    if chunk_size == -1:
        chunk_limit = -1

    # Get largest segment bytes
    pattern_largest_segment = list(filter(None, pattern.split('??')))
    pattern_largest_segment.sort(key=len, reverse=True)
    pattern_largest_segment = pattern_largest_segment[0]
    pattern_largest_segment_position = pattern.index(
        pattern_largest_segment) // 2
    pattern_largest_segment = bytes.fromhex(pattern_largest_segment)

    # Search method 1 (no wildcards)
    if pattern.count('?') == 0:
        pattern_bytes = bytes.fromhex(pattern)
        chunk_position = 0
        with open(path, 'rb') as f:
            if start_offset > 0:
                f.seek(start_offset)
            while True:
                if chunk_limit > 0:
                    if chunk_position / chunk_size >= chunk_limit:
                        return
                try:
                    data = f.read(chunk_size)
                except MemoryError:
                    yield False
                if not data:
                    return
                i = 0
                found_position = 0
                while True:
                    try:
                        found_position = data.index(pattern_bytes,
                                                    found_position + i)
                        if chunk_size > 0:
                            yield chunk_position + found_position + \
                                  start_offset
                        else:
                            yield found_position + start_offset
                        if find_all is False:
                            return
                    except ValueError:
                        break
                    i += 1
                chunk_position += chunk_size
                continue
            return

    # Create a list of wildcard positions
    pattern_wildcard_positions = []
    for i in range(0, len(pattern), 2):
        pattern_byte = pattern[i:i + 2]
        if pattern_byte == '??':
            pattern_wildcard_positions.append(i // 2)

    # Remove wildcards from pattern string and convert to bytes
    pattern_len = len(pattern) // 2
    pattern_bytes = pattern.replace('?', '')
    pattern_bytes = bytes.fromhex(pattern_bytes)

    # Search method 2 (wildcards)
    possible_positions = []
    end_of_file = False
    first_result = True
    chunk_position = 0

    with open(path, 'rb') as f:
        if start_offset > 0:
            f.seek(start_offset)
        while not end_of_file:
            if chunk_limit > 0:
                if chunk_position / chunk_size >= chunk_limit:
                    return
            try:
                data = f.read(chunk_size)
            except MemoryError:
                yield False
            if not data:
                end_of_file = True
            chunk_search = True
            while chunk_search:
                try:
                    if first_result is True:
                        possible_positions.append(
                            data.index(pattern_largest_segment))
                        first_result = False
                    else:
                        possible_positions.append(
                            data.index(pattern_largest_segment,
                                       possible_positions[-1] + 1))
                except ValueError:
                    if chunk_size > 0:
                        chunk_position += chunk_size
                    chunk_search = False
            for possible_position in possible_positions:
                possible_position -= pattern_largest_segment_position
                match_count = 0
                pattern_bytes_pos = 0
                data_offset_pos = 0
                i = 0
                while i < pattern_len:
                    if i in pattern_wildcard_positions:
                        match_count += 1
                        data_offset_pos += 1
                        i += 1
                        continue
                    elif pattern_bytes[pattern_bytes_pos] == data[
                                possible_position + data_offset_pos]:
                        match_count += 1
                        data_offset_pos += 1
                        pattern_bytes_pos += 1
                        i += 1
                        continue
                    i += 1
                if match_count == pattern_len:
                    if find_all is True:
                        if chunk_size > 0:
                            yield chunk_position + possible_position + \
                                  start_offset - chunk_size
                        else:
                            yield possible_position + start_offset
                    else:
                        yield possible_position + chunk_position + \
                              start_offset - chunk_size
                        return
            possible_positions = []
            first_result = True
        return
