#!/usr/bin/env sh

MYDIR="$(dirname ${0})"
FN_FILE="${1:-${MYDIR}/ntdll_zw_functions.txt}"

TYPEDEFS=""
STATICS=""
CURLINE=0
while read -r line; do
    CURLINE=$(expr ${CURLINE} + 1)
    VALID=1

    rtype=$(printf '%s\n' "${line}" | grep -oE '(NTSTATUS NTAPI|VOID NTAPI|PVOID NTAPI)')
    if [ -z "${rtype}" ]; then
        printf '%s\n' "Line ${CURLINE}: Missing return value of either type 'NTSTATUS NTAPI' or 'VOID NTAPI'." >&2
        VALID=0
    fi

    fnname=$(printf '%s\n' "${line}" | grep -oE '(Zw|Rtl|Ob|Mm|Io)[^ (]*')
    if [ -z "${fnname}" ]; then
        printf '%s\n' "Line ${CURLINE}: Missing function name." >&2
        VALID=0
    fi

    fnsig=$(printf '%s\n' "${line}" | grep -oE '\([^;]*')
    if [ -z "${fnsig}" ]; then
        printf '%s\n' "Line ${CURLINE}: Missing function signature." >&2
        VALID=0
    fi

    params_without_braces=$(printf '%s\n' "${fnsig}" | tr -d '()')
    if [ ! -z "${params_without_braces}" ]; then
        param_names=$(printf '%s\n' "${params_without_braces}" | sed 's/\([^,]*\)/\1\n/g' | grep -oE '[^ ]*$')
        if [ -z "${param_names}" ]; then
            printf '%s\n' "Line ${CURLINE}: Could not parse function parameters." >&2
            VALID=0
        fi
    else
        param_names=""
    fi
    params=""
    for param in ${param_names}; do
        if [ -z "${param}" ]; then
            printf '%s\n' "Line ${CURLINE}: Invalid parameter found. Please re-check regex'es used." >&2
            VALID=0
        fi
        params="${params}${param}, "
    done
    params=$(printf '%s\n' "${params}" | sed 's/^\(.*\), $/\1/g')
    if [ -z "${params}" -a ! -z "${params_without_braces}" ]; then
        printf '%s\n' "Line ${CURLINE}: Parameters empty. Please re-check regex'es used." >&2
        VALID=0
    fi

    if [ ${VALID} -eq 1 ]; then
        TYPEDEFS="${TYPEDEFS}\ntypedef ${rtype} (*${fnname}_t) ${fnsig};"
        STATICS="${STATICS}\nstatic ${fnname}_t _${fnname} = NULL;"
        INITS=$(cat <<EOF
${INITS}
#ifdef __cplusplus
    RtlInitUnicodeString(&fnName, skCrypt(L"${fnname}"));
#else
    RtlInitUnicodeString(&fnName, L"${fnname}");
#endif
    _${fnname} = (${fnname}_t)MmGetSystemRoutineAddress(&fnName);
    if (_${fnname} == NULL)
    {
#ifdef __cplusplus
        DbgPrint(skCrypt("%s\\\n"), skCrypt("System routine ${fnname} not found."));
#else
        DbgPrint("%s\\\n", "System routine ${fnname} not found.");
#endif
        retval++;
    }
EOF
        )
        WRAPPERS=$(cat <<EOF
${WRAPPERS}

${rtype} ${fnname} ${fnsig}
{
EOF
        )
        case $rtype in
        NTSTATUS*)
            WRAPPERS=$(cat <<EOF
${WRAPPERS}
    if (_${fnname} == NULL)
        return STATUS_PROCEDURE_NOT_FOUND;

    return _${fnname} (${params});
}

${rtype} Wrapper${fnname} ${fnsig}
{
    return _${fnname} (${params});
}
EOF
            )
        ;;
        PVOID*)
            WRAPPERS=$(cat <<EOF
${WRAPPERS}
    return _${fnname} (${params});
}

${rtype} Wrapper${fnname} ${fnsig}
{
    return _${fnname} (${params});
}
EOF
            )
        esac
    fi
done < "${FN_FILE}"

cat <<EOF
/* This file was auto generated by $(basename ${0}) */
#include <ntddk.h>

#ifdef __cplusplus
#define _KERNEL_MODE 1
#include "obfuscate.hpp"

extern "C" {
#endif
EOF
echo "${TYPEDEFS}"
echo "${STATICS}"
cat <<EOF

int __cdecl $(basename -a -s '.txt' ${FN_FILE}) (void)
{
    int retval = 0;
    UNICODE_STRING fnName;
EOF
echo "${INITS}"
cat <<EOF

    return retval;
}
EOF
echo "${WRAPPERS}"

cat <<EOF

#ifdef __cplusplus
};
#endif
EOF