import os
import sys
import random
import string
import shutil
import time
import paramiko
import gateway
import filecmp
from scp import SCPClient

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(SCRIPT_DIR, '..', 'Data')
FILES_DIR = os.path.join(SCRIPT_DIR, 'Files')

def generate_random_string(length, chars):
    return ''.join(random.choice(chars) for _ in range(length))

def get_hex_offset(filename, offset, length):
    with open(filename, 'rb') as f:
        f.seek(offset)
        return f.read(length)

def get_mtd_partitions(client):
    stdin, stdout, stderr = client.exec_command('cat /proc/mtd')
    data = stdout.read().decode()
    lines = data.split('\n')
    partitions = {}
    for line in lines:
        if 'u-boot' in line:
            partitions['u-boot'] = line.split(':')[0]
        elif 'radio' in line:
            partitions['radio'] = line.split(':')[0]
        elif 'config' in line:
            partitions['config'] = line.split(':')[0]
    return partitions

def increment_mac_address(mac):
    """Увеличивает MAC-адрес на 1"""
    mac_int = int.from_bytes(mac, byteorder='big')
    mac_int += 1
    return mac_int.to_bytes(6, byteorder='big')

def modify_file(input_file_path, output_file_path, base_mac):
    with open(input_file_path, "rb") as f:
        data = bytearray(f.read())

    offsets = [0x4, 0x28, 0x404, 0x8004]
    
    # Устанавливаем базовый MAC-адрес на смещение 0x28
    data[0x28:0x28+6] = base_mac
    base_mac = increment_mac_address(base_mac)  # Инкрементируем базовый MAC

    # Устанавливаем инкрементированные MAC-адреса на остальные смещения
    for offset in offsets:
        if offset != 0x28:  # Пропускаем смещение 0x28
            data[offset:offset+6] = base_mac
            base_mac = increment_mac_address(base_mac)

    with open(output_file_path, "wb") as f:
        f.write(data)

def replace_values(filename):
    file_path = os.path.join(FILES_DIR, filename)

    try:
        with open(file_path, 'rb') as f:
            data = f.read()
    except FileNotFoundError:
        print(f'Файл {file_path} не найден. Пожалуйста, поместите ваш u-config.bin в директорию Files.')
        return

    patterns = {
        'servicetag': (b'\x73\x65\x72\x76\x69\x63\x65\x74\x61\x67\x3D', string.digits),
        'sernumb': (b'\x73\x65\x72\x6E\x75\x6D\x62\x3D', string.digits),
        'servicepass': (b'\x73\x65\x72\x76\x69\x63\x65\x70\x61\x73\x73\x3D', string.ascii_letters + string.digits)
    }

    for name, (pattern, chars) in patterns.items():
        start = data.find(pattern)
        if start != -1:
            start += len(pattern)
            end = data.find(b'\x00', start)
            if end != -1:
                if name == 'sernumb': 
                    start = end - 4
                new_value = generate_random_string(end - start, chars).encode()
                data = data[:start] + new_value + data[end:]
                # print(f'Значение переменной {name} было успешно заменено.')
            else:
                print(f'Не найдено "00" после переменной {name}.')
        else:
            print(f'Переменная {name} не найдена.')

    with open(file_path, 'wb') as f:
        f.write(data)
        # print(f'Новые данные были успешно записаны в исходный файл {file_path}')

def find_and_replace(data, search_bytes, replace_file, length):
    replace_file_path = os.path.join(FILES_DIR, replace_file)  # Изменить путь на папку Files

    start_index = 0
    while True:
        start_index = data.find(search_bytes, start_index)
        if start_index == -1:
            break
        # print(f"Найдено смещение {search_bytes.hex()}. Заменяю данными из {replace_file}...")
        with open(replace_file_path, 'rb') as rf:
            replace_data = rf.read()
            if len(replace_data) != length:
                print(f"Предупреждение: {replace_file} не имеет ожидаемой длины")
            data = data[:start_index] + replace_data[:length] + data[start_index+length:]
            # print(f"Замена завершена.")
        start_index += length
    return data


def backup(client, partitions):
    commands = [
        (f'dd if=/dev/{partitions["u-boot"]} of=/tmp/u-boot_stock.bin count=512', 'u-boot_stock.bin'),
        (f'dd if=/dev/{partitions["radio"]} of=/tmp/eeprom_stock.bin count=128', 'eeprom_stock.bin'),
        (f'dd if=/dev/{partitions["config"]} of=/tmp/config_stock.bin count=128', 'config_stock.bin')
    ]
    for command, filename in commands:
        stdin, stdout, stderr = client.exec_command(command, get_pty=True)
        time.sleep(2)
        data = stdout.read() + stderr.read()
        scp = SCPClient(client.get_transport())
        scp.get(f'/tmp/{filename}', os.path.join(DATA_DIR, filename))

        if filename == 'config_stock.bin':
            # Try to get base MAC address from the config_stock.bin file
            config_path = os.path.join(DATA_DIR, filename)
            if os.path.exists(config_path):
                base_mac = get_hex_offset(config_path, 0x8, 6)
                os.remove(config_path)  # Remove the config_stock.bin after using it
            else:
                base_mac = None

        if filename == 'u-boot_stock.bin':
            print(f'Стоковый загрузчик {filename} успешно сохранен в папку Data')
        elif filename == 'eeprom_stock.bin':
            print(f'{filename} успешно сохранен в папку Data')
        else:
            print(f'Ошибка при сохранении {filename}')

    # Make sure base_mac is not None before proceeding
    if base_mac:
        # Modify EEPROM file
        modify_file(os.path.join(DATA_DIR, 'eeprom_stock.bin'), os.path.join(FILES_DIR, "eeprom_keenetic.bin"), base_mac)

        # Modify u-config file
        replace_values('u-config.bin')

        # Create full.bin
        create_full_bin(base_mac.hex().upper())
    else:
        print("Не удалось получить MAC-адрес из config_stock.bin. Проверьте файл.")

def write_loader(client):
    stdin, stdout, stderr = client.exec_command('insmod mtd-rw i_want_a_brick=1')
    scp = SCPClient(client.get_transport())
    scp.put(os.path.join(FILES_DIR, 'Breed_C6U.bin'), '/tmp/Breed_C6U.bin')
    stdin, stdout, stderr = client.exec_command('mtd write /tmp/Breed_C6U.bin u-boot', get_pty=True)
    time.sleep(2)
    data = stdout.read() + stderr.read()

def create_full_bin(base_mac):
    search_bytes_uconfig = bytes.fromhex("73 3D 6C 1E 62 6F 6F 74 63")
    search_bytes_eeprom = bytes.fromhex("03 76 0B 00")
    length = 0x10000

    print("Собираю готовый образ...")
    with open(os.path.join(FILES_DIR, 'full.bin'), 'rb') as f:
        data = f.read()

    data = find_and_replace(data, search_bytes_uconfig, 'u-config.bin', length)
    data = find_and_replace(data, search_bytes_eeprom, 'eeprom_keenetic.bin', length)

    mac_address_offset = 0x40028
    mac_address_length = 6
    mac_address = data[mac_address_offset:mac_address_offset+mac_address_length]
    mac_address_hex = mac_address.hex().upper()

    # Determine the target directory
    current_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(current_dir)
    grandparent_dir = os.path.dirname(parent_dir)
    target_dir = os.path.join(grandparent_dir, 'Keenetic')
    
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)

    print(f"Сохраняю измененные данные в файл Keenetic-C6U-{mac_address_hex}.bin...")
    with open(os.path.join(target_dir, f"Keenetic-C6U-{mac_address_hex}.bin"), 'wb') as f:
        f.write(data)
    print("Сохранение завершено, файл создан в папке Keenetic")

def main():
    print ("+----------------------------------------------------------------+")
    print ("|        Breed installer for TP-Link C6Uv1 by spatiumstas        |")
    print ("+----------------------------------------------------------------+")

    router_ip = gateway.get_ip_address()
    username = 'root'
    password = 'root'
    ssh_port = '22'

    try:
        with paramiko.SSHClient() as client:
            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            client.connect(hostname=router_ip, username=username, password=password, port=ssh_port)
            time.sleep(2)
            partitions = get_mtd_partitions(client)
            backup(client, partitions)
            write_loader(client)
            print ("")
            print("+--------------------------------------------------------------+")
            print("|              Загрузчик Breed успешно установлен              |")
            print("+--------------------------------------------------------------+")
            print ("")
            print ('Перезагрузка в Breed...')
            client.exec_command('reboot')
            time.sleep(2)
            client.close()
            sys.exit()
    except Exception as e:
        print(f"Произошла ошибка: {e}")
    finally:
        client.close()

if __name__ == "__main__":
    main()
