import sys
import struct
import argparse

def is_big_endian_elf(file_path):
    with open(file_path, 'rb') as f:
        f.seek(0)  # ELF header starts at the beginning
        elf_header = f.read(16)  # ELF header is at least 16 bytes
        # ELF magic number check (first 4 bytes should be 0x7F 'E' 'L' 'F')
        if elf_header[:4] != b'\x7fELF':
            raise ValueError("Not a valid ELF file")
        # Byte 5 in the ELF header specifies endianness
        # 1 means little-endian, 2 means big-endian
        ei_data = elf_header[5]
        return ei_data == 2

def find_patterns(file_path, old_pattern):
    with open(file_path, 'rb') as f:
        file_content = f.read()
    
    indices = []
    start = 0
    while True:
        index = file_content.find(old_pattern, start)
        if index == -1:
            break
        indices.append(index)
        start = index + len(old_pattern)
    
    return indices

def replace_patterns(file_path, output_path, indices, old_pattern, new_pattern):
    with open(file_path, 'rb') as f:
        file_content = f.read()
    
    new_file_content = bytearray(file_content)
    
    for index in indices:
        # Ensure the index is aligned on a 4-byte boundary
        if index % 4 != 0:
            raise ValueError(f"Pattern found at non-aligned position: 0x{index:08X}. SPARC requires 4-byte alignment.")
        
        # Replace the old pattern with the new one
        new_file_content[index:index + len(old_pattern)] = new_pattern
    
    with open(output_path, 'wb') as f:
        f.write(new_file_content)
    
    print(f"Pattern replaced successfully. Output written to {output_path}")

def main():
    parser = argparse.ArgumentParser(description="Find and replace binary patterns in an ELF file.")
    parser.add_argument('input_file', help='Input ELF file path')
    parser.add_argument('output_file', help='Output file path')
    parser.add_argument('--yes', action='store_true', help='Automatically answer "yes" to patching prompt')
    
    args = parser.parse_args()

    input_file = args.input_file
    output_file = args.output_file
    auto_yes = args.yes

    # Define the old and new patterns as raw bytes (including NOP in the new pattern)
    old_pattern = b'\x9d\xe3\xbe\x68\x92\x0e\xa0\xff\x90\x10\x20\x01\xa2\x0e\xe0\xff\xa6\x10\x00\x1c\xf0\x23\xa0\x78'
    new_pattern = b'\x9d\xe3\xbe\x68\x92\x0e\xa0\xff\x90\x10\x20\x00\x81\xc7\xe0\x08\x01\x00\x00\x00\x01\x00\x00\x00'
    
    # Step 1: Check if the input file is a valid ELF file and determine endianness
    try:
        if is_big_endian_elf(input_file):
            print("The ELF file is big-endian (SPARC compatible).")
        else:
            print("Warning: The ELF file is little-endian. This might not be compatible with SPARC.")
            sys.exit(1)
    except ValueError as e:
        print(e)
        sys.exit(1)
    
    # Step 2: Analyze the file and find pattern locations
    print("Analysis Step:")
    indices = find_patterns(input_file, old_pattern)
    
    if not indices:
        print("Pattern not found in the file.")
        sys.exit(0)
    
    print(f"Pattern found {len(indices)} time(s) in {input_file} at the following positions:")
    for index in indices:
        print(f"0x{index:08X}")
    
    # Step 3: Ask the user if they want to proceed with the replacement (unless --yes is passed)
    if auto_yes:
        user_input = 'yes'
    else:
        user_input = input("Do you want to apply the replacement? (yes/no): ").strip().lower()
    
    if user_input == 'yes':
        replace_patterns(input_file, output_file, indices, old_pattern, new_pattern)
    else:
        print("No changes made. Exiting.")

if __name__ == "__main__":
    main()