import subprocess
import os
import re

import ogGlobals
import SystemLib
import FileLib

def chntpw(*args):
    chntpw_path = subprocess.check_output(['which', 'drbl-chntpw']).decode().strip()
    if not chntpw_path:
        chntpw_path = subprocess.check_output(['which', 'chntpw']).decode().strip()
    subprocess.run([chntpw_path, '-e'] + list(args), timeout=5)

def ogAddRegistryKey(path_mountpoint, str_hive, str_key):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Añadir nueva clave.
    chntpw(FILE, 'cd', str_key.rsplit('\\', 1)[0])
    chntpw(FILE, 'nk', str_key.rsplit('\\', 1)[-1])
    chntpw(FILE, 'q')
    chntpw(FILE, 'y')

def ogAddRegistryValue(path_mountpoint, str_hive, str_key, str_valuename, str_valuetype=''):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Determine the value type.
    if str_valuetype.upper() == 'STRING' or str_valuetype == '':
        TYPE = 1
    elif str_valuetype.upper() == 'BINARY':
        TYPE = 3
    elif str_valuetype.upper() == 'DWORD':
        TYPE = 4
    else:
        SystemLib.ogRaiseError(OG_ERR_OUTOFLIMIT, str_valuetype)
        return

    # Add the registry value.
    chntpw(FILE, 'cd', str_key.rsplit('\\', 1)[0])
    chntpw(FILE, 'nv', str_valuename.rsplit('\\', 1)[-1], TYPE)
    chntpw(FILE, 'q')
    chntpw(FILE, 'y')

def ogDeleteRegistryKey(path_mountpoint, str_hive, str_key):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Delete the registry key.
    subprocess.run(['chntpw', FILE], input=f"cd {str_key.rsplit('\\', 1)[0]}\ndk {str_key.rsplit('\\', 1)[-1]}\nq\ny\n".encode(), timeout=5)

def ogDeleteRegistryValue(path_mountpoint, str_hive, str_valuename):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Delete the registry value.
    chntpw(FILE, 'cd', str_valuename.rsplit('\\', 1)[0])
    chntpw(FILE, 'dv', str_valuename.rsplit('\\', 1)[-1])
    chntpw(FILE, 'q')
    chntpw(FILE, 'y')

def ogGetHivePath(path_mountpoint, str_hive):
    # Variables locales.
    FILE = None

    # Camino del fichero de registro de usuario o de sistema (de menor a mayor prioridad).
    FILE = FileLib.ogGetPath(f"/{path_mountpoint}/Windows/System32/config/{str_hive}")
    if not FILE:
        FILE = FileLib.ogGetPath(f"/{path_mountpoint}/Users/{str_hive}/NTUSER.DAT")
    if not FILE:
        FILE = FileLib.ogGetPath(f"/{path_mountpoint}/winnt/system32/config/{str_hive}")
    if not FILE:
        FILE = FileLib.ogGetPath(f"/{path_mountpoint}/Documents and Settings/{str_hive}/NTUSER.DAT")
    if FILE and os.path.isfile(FILE):
        return FILE
    else:
        SystemLib.ogRaiseError(OG_ERR_NOTFOUND, f"{path_mountpoint} {str_hive}")
        return None
    
def ogGetRegistryValue(path_mountpoint, str_hive, str_valuename):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Devolver el dato del valor de registro.
    chntpw_cmd = f'''
        chntpw "{FILE}" << EOT 2> /dev/null | awk '/> Value/ {{
            if (index($0, "REG_BINARY") > 0) {{
                data=""
            }} else {{
                getline
                data=$0
            }}
        }}
        /^:[0-9A-F]+ / {{
            data=data""substr($0, 9, 48)
        }}
        END {{
            print data
        }}'
        cd {str_valuename.rsplit('\\', 1)[0]}
        cat {str_valuename.rsplit('\\', 1)[-1]}
        q
        EOT
    '''
    subprocess.run(chntpw_cmd, shell=True, timeout=5)
    
def ogListRegistryKeys(path_mountpoint, str_hive, str_key):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Devolver la lista de claves de registro.
    chntpw_cmd = f'''
        chntpw "{FILE}" << EOT 2> /dev/null | awk 'BEGIN {{FS="[<>]"}} $1~/^  $/ {{print $2}}'
        ls {str_key}
        q
        EOT
    '''
    subprocess.run(chntpw_cmd, shell=True, timeout=5)

def ogListRegistryValues(path_mountpoint, str_hive, str_key):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Devolver la lista de valores de registro.
    chntpw_cmd = f'''
        chntpw "{FILE}" << EOT 2> /dev/null | awk 'BEGIN {{FS="[<>]"}} $1~/REG_/ {{print $2}}'
        ls {str_key}
        q
        EOT
    '''
    subprocess.run(chntpw_cmd, shell=True, timeout=5)

def ogSetRegistryValue(path_mountpoint, str_hive, str_valuename, str_data):
    # Variables locales.
    FILE = ogGetHivePath(path_mountpoint, str_hive)
    if not FILE:
        return

    # Fichero temporal para componer la entrada al comando "chntpw".
    tmpfile = "/tmp/chntpw$$"
    try:
        # Comprobar tipo de datos del valor del registro.
        with open(tmpfile, 'w') as f:
            f.write(f"ls {str_valuename.rsplit('\\', 1)[0]}\nq\n")
        output = subprocess.check_output(['chntpw', FILE], input=open(tmpfile, 'rb'), stderr=subprocess.DEVNULL).decode()
        if f"BINARY.*<{str_valuename.rsplit('\\', 1)[-1]}>" in output:
            # Procesar tipo binario (incluir nº de bytes y líneas de 16 parejas hexadecimales).
            if not re.match(r'^([0-9A-F]{2} )*$', str_data):
                SystemLib.ogRaiseError(OG_ERR_FORMAT, f'"{str_data}"')
                return
            n = len(str_data) + 1
            with open(tmpfile, 'w') as f:
                f.write(f"cd {str_valuename.rsplit('\\', 1)[0]}\ned {str_valuename.rsplit('\\', 1)[-1]}\n{int(n/3)}\n")
                for i in range(0, n, 48):
                    f.write(f":{i//3:05x}  {str_data[i:i+48]}\n")
                f.write("s\nq\ny\n")
        else:
            # Cambiar el dato del valor de registro para cadenas y bytes.
            with open(tmpfile, 'w') as f:
                f.write(f"cd {str_valuename.rsplit('\\', 1)[0]}\ned {str_valuename.rsplit('\\', 1)[-1]}\n{str_data}\nq\ny\n")
        
        # Aplicar cambios.
        subprocess.run(['chntpw', FILE], input=open(tmpfile, 'rb'), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    finally:
        os.remove(tmpfile)