#/**
#@file    RegistryLib.py
#@brief   Librería o clase Registry
#@class   Boot
#@brief   Funciones para gestión del registro de Windows.
#@warning License: GNU GPLv3+
#*/

import subprocess
import os
import re
import shutil
import tempfile

import ogGlobals
import SystemLib
import FileLib

# Función ficticia para lanzar chntpw con timeout de 5 s., evitando cuelgues del programa.
chntpw_exe = shutil.which ('drbl-chntpw') or shutil.which ('chntpw')
def chntpw (hivefile, input_file):
    with open (input_file, 'r') as fd:
        input_contents = fd.read()
    return subprocess.run ([chntpw_exe, '-e', hivefile], timeout=5, input=input_contents, capture_output=True, text=True).stdout

## en el codigo bash aparecen "${3%\\*}" y "${3##*\\}" varias veces
## ${3%\\*}   es el "dirname" de una key del registro
## ${3##*\\}  es el "basename"
def _split_k (k):
    k_elems = k.split ('\\')
    k_dirname = '\\'.join (k_elems[0:-1])
    k_basename = k_elems[-1]
    return k_dirname, k_basename

#/**
#         ogAddRegistryKey path_mountpoint str_hive str_keyname
#@brief   Añade una nueva clave al registro de Windows.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_keyname      nombre de la clave
#@return  (nada)
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { default, sam, security, software, system, components }
#@warning Requisitos: chntpw
#@warning El sistema de archivos de Windows debe estar montada previamente.
#*/ ##
def ogAddRegistryKey (mntpt, hive, k):
    hivefile = ogGetHivePath (mntpt, hive)
    if not hivefile: return

    k_dirname, k_basename = _split_k (k)

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'cd {k_dirname}\n')
        f.write (f'nk {k_basename}\n')
        f.write ('q\ny\n')
        f.close()
        chntpw (hivefile, f.name)
        os.remove (f.name)

#/**
#         ogAddRegistryValue path_mountpoint str_hive str_valuename [str_valuetype]
#@brief   Añade un nuevo valor al registro de Windows, indicando su tipo de datos.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_valuename    nombre del valor
#@param   str_valuetype    tipo de datos del valor (opcional)
#@return  (nada)
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { DEFAULT, SAM, SECURITY, SOFTWARE, SYSTEM, COMPONENTS }
#@note    valuetype = { STRING, BINARY, DWORD }, por defecto: STRING
#@warning Requisitos: chntpw
#@warning El sistema de archivos de Windows debe estar montada previamente.
#*/ ##
#ogAddRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Microsoft\NewKey\Value1')    ## type STRING by default
#ogAddRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Microsoft\NewKey\Value1', 'STRING')
#ogAddRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Microsoft\NewKey\Value1', 'BINARY')
#ogAddRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Microsoft\NewKey\Value1', 'DWORD')
def ogAddRegistryValue (mntpt, hive, k, vtype='STRING'):
    hivefile = ogGetHivePath (mntpt, hive)
    if not hivefile: return

    k_dirname, k_basename = _split_k (k)

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

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'cd {k_dirname}\n')
        f.write (f'nv {TYPE} {k_basename}\n')
        f.write ('q\ny\n')
        f.close()
        chntpw (hivefile, f.name)
        os.remove (f.name)


#/**
#         ogDeleteRegistryKey path_mountpoint str_hive str_keyname
#@brief   Elimina una clave del registro de Windows con todo su contenido.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_keyname      nombre de la clave
#@return  (nada)
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { default, sam, security, software, system, components }
#@warning Requisitos: chntpw
#@warning El sistema de archivos de Windows debe estar montada previamente.
#@warning La clave debe estar vacía para poder ser borrada.
#*/ ##
def ogDeleteRegistryKey (mntpt, hive, k):
    hivefile = ogGetHivePath (mntpt, hive)
    if not hivefile: return

    k_dirname, k_basename = _split_k (k)

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'cd {k_dirname}\n')
        f.write (f'dk {k_basename}\n')
        f.write ('q\ny\n')
        f.close()
        chntpw (hivefile, f.name)
        os.remove (f.name)


#/**
#         ogDeleteRegistryValue path_mountpoint str_hive str_valuename
#@brief   Elimina un valor del registro de Windows.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_valuename    nombre del valor
#@return  (nada)
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { default, sam, security, software, system, components }
#@warning Requisitos: chntpw
#@warning El sistema de archivos de Windows debe estar montada previamente.
#*/ ##
#ogDeleteRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Microsoft\NewKey\Value1')
def ogDeleteRegistryValue (mntpt, hive, k):
    hivefile = ogGetHivePath (mntpt, hive)
    if not hivefile: return

    k_dirname, k_basename = _split_k (k)

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'cd {k_dirname}\n')
        f.write (f'dv {k_basename}\n')
        f.write ('q\ny\n')
        f.close()
        chntpw (hivefile, f.name)
        os.remove(f.name)


#/**
#         ogGetHivePath path_mountpoint [str_hive|str_user]
#@brief   Función básica que devuelve el camino del fichero con una sección del registro.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@return  str_path - camino del fichero de registro
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { DEFAULT, SAM, SECURITY, SOFTWARE, SYSTEM, COMPONENTS, NombreDeUsuario }
#@warning El sistema de archivos de Windows debe estar montada previamente.
#*/ ##
#ogGetHivePath ('/mnt/sda1', 'user1')     =>  /mnt/sda1/Users/user1/NTUSER.DAT
#ogGetHivePath ('/mnt/sda1', 'SYSTEM')    =>  //mnt/sda1/Windows/System32/config/SYSTEM
#ogGetHivePath ('/mnt/sda1', 'IEUser')    =>  //mnt/sda1/Users/IEUser/NTUSER.DAT
def ogGetHivePath(mntpt, hive):
    # Camino del fichero de registro de usuario o de sistema (de menor a mayor prioridad).
    FILE              = FileLib.ogGetPath(file=f"/{mntpt}/Windows/System32/config/{hive}")
    if not FILE: FILE = FileLib.ogGetPath(file=f"/{mntpt}/Users/{hive}/NTUSER.DAT")
    if not FILE: FILE = FileLib.ogGetPath(file=f"/{mntpt}/winnt/system32/config/{hive}")
    if not FILE: FILE = FileLib.ogGetPath(file=f"/{mntpt}/Documents and Settings/{hive}/NTUSER.DAT")
    if FILE and os.path.isfile(FILE):
        return FILE
    else:
        SystemLib.ogRaiseError ([], ogGlobals.OG_ERR_NOTFOUND, f'{mntpt} {hive}')
        return None


## simulate 'grep --after-context 1'
def _grep_A1 (strings, search_term):
    results = []
    for i in range (len (strings)):
        if search_term in strings[i]:
            results.append (strings[i])
            if i + 1 < len(strings):
                results.append (strings[i + 1])

    return results

#/**
#         ogGetRegistryValue path_mountpoint str_hive str_valuename
#@brief   Devuelve el dato de un valor del registro de Windows.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_valuename    nombre del valor
#@return  str_valuedata - datos del valor.
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { default, sam, security, software, system, components }
#@warning Requisitos: chntpw, awk
#@warning El sistema de archivos de Windows debe estar montado previamente.
#*/ ##
def ogGetRegistryValue (mntpt, hive, k):
    hivefile = ogGetHivePath(mntpt, hive)
    if not hivefile: return

    k_dirname, k_basename = _split_k (k)

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'cd {k_dirname}\n')
        f.write (f'cat {k_basename}\n')
        f.write ('q\n')
        f.close()
        chntpw_out = chntpw (hivefile, f.name)
        os.remove (f.name)

    lines = chntpw_out.splitlines()
    lines = _grep_A1 (lines, '> Value')
    if 2 != len (lines):
        return None

    ret = None
    if 'REG_BINARY' in lines[0]:
        if re.search ('^:[0-9A-F]+ ', lines[1]):
            ret = lines[1][8:56]
    else:
        ret = lines[1]
    return ret


#/**
#         ogListRegistryKeys path_mountpoint str_hive str_key
#@brief   Lista los nombres de subclaves de una determinada clave del registro de Windows.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_key          clave de registro
#@return  str_subkey ... - lista de subclaves
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { default, sam, security, software, system, components }
#@warning Requisitos: chntpw, awk
#@warning El sistema de archivos de Windows debe estar montado previamente.
#*/ ##
#ogListRegistryKeys ('/mnt/sda1', 'SOFTWARE', '\Microsoft\Windows\CurrentVersion')
def ogListRegistryKeys (mntpt, hive, k):
    hivefile = ogGetHivePath(mntpt, hive)
    if not hivefile: return

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'ls {k}\n')
        f.write ('q\n')
        f.close()
        chntpw_out = chntpw (hivefile, f.name)
        os.remove (f.name)

    lines = chntpw_out.splitlines()
    ret = []
    for l in lines:
        elems = re.split ('[<>]', l)
        if len(elems) < 2: continue
        if '  ' == elems[0]:
            ret.append (elems[1])
    return ret


#/**
#         ogListRegistryValues path_mountpoint str_hive str_key
#@brief   Lista los nombres de valores de una determinada clave del registro de Windows.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_key          clave de registro
#@return  str_value ... - lista de valores
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { default, sam, security, software, system, components }
#@warning Requisitos: chntpw, awk
#@warning El sistema de archivos de Windows debe estar montado previamente.
#*/ ##
#ogListRegistryValues ('/mnt/sda1', 'SOFTWARE', '\Microsoft\Windows\CurrentVersion')
def ogListRegistryValues (mntpt, hive, k):
    hivefile = ogGetHivePath(mntpt, hive)
    if not hivefile: return

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'ls {k}\n')
        f.write ('q\n')
        f.close()
        chntpw_out = chntpw (hivefile, f.name)
        os.remove (f.name)

    lines = chntpw_out.splitlines()
    ret = []
    for l in lines:
        elems = re.split ('[<>]', l)
        if len(elems) < 2: continue
        if 'REG_' in elems[0]:
            ret.append (elems[1])
    return ret


def _format_hex (hex_string):
    result = []
    offset = 0

    hex_values = hex_string.strip().split()
    result.append (str (len (hex_values)))

    while hex_values:
        chunk = hex_values[:16]
        hex_values = hex_values[16:]

        offset_str = f':{offset:05x}  '
        hex_line = ' '.join (chunk)
        result.append (offset_str + hex_line)

        offset += 16

    return '\n'.join (result)

#/**
#         ogSetRegistryValue path_mountpoint str_hive str_valuename str_valuedata
#@brief   Establece el dato asociado a un valor del registro de Windows.
#@param   path_mountpoint  directorio donde está montado el sistema Windows
#@param   str_hive         sección del registro
#@param   str_valuename    nombre del valor de registro
#@param   str_valuedata    dato del valor de registro
#@return  (nada)
#@exception OG_ERR_FORMAT    Formato incorrecto.
#@exception OG_ERR_NOTFOUND  Fichero de registro no encontrado.
#@note    hive = { default, sam, security, software, system, components }
#@warning Requisitos: chntpw
#@warning El sistema de archivos de Windows debe estar montado previamente.
#*/ ##
#ogSetRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Key\SubKey\StringValue', 'Abcde Fghij')
#ogSetRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Key\SubKey\DwordValue', 1)
#ogSetRegistryValue ('/mnt/sda1', 'SOFTWARE', '\Key\SubKey\BinaryValue', '04 08 0C 10')
def ogSetRegistryValue (mntpt, hive, k, v):
    hivefile = ogGetHivePath (mntpt, hive)
    if not hivefile: return

    k_dirname, k_basename = _split_k (k)

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f"ls {k_dirname}\n")
        f.write ('q\n')
        f.close()
        chntpw_out = chntpw (hivefile, f.name)
        os.remove(f.name)

    if re.search (f"BINARY.*<{k_basename}>", chntpw_out):
        ## the entry in the registry is binary. Our input should be a sequence of bytes

        if ' ' != v[-1]: v += ' '     ## the regex below requires a trailing space
        if not re.match (r'^([0-9A-F]{2} )*$', v.upper()):
            SystemLib.ogRaiseError ([], ogGlobals.OG_ERR_FORMAT, f'"{v}"')
            return

        formatted = _format_hex (v.upper())
        formatted += '\ns'
    else:
        formatted = v

    with tempfile.NamedTemporaryFile (delete_on_close=False, prefix='chntpw-', mode='w') as f:
        f.write (f'cd {k_dirname}\n')
        f.write (f'ed {k_basename}\n')
        f.write (f'{formatted}\n')
        f.write ('q\ny\n')
        f.close()
        chntpw (hivefile, f.name)
        os.remove(f.name)