/*
    Copyright (C) 2003 Olaf Flebbe, Science and Computing AG
                       o.flebbe@science-computing.de
    Copyright (C) 2013 Timothy Pearson, Northern Illinois University
	                   kb9vqf@pearsoncomputing.net

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.

*/
#define UNICODE
#define SECURITY_WIN32
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <ntsecapi.h>
#include <sspi.h>
#include <ntsecpkg.h>
#include <stdio.h>
//#include <ntstatus.h>
#include <string.h>
//#include <subauth.h>
#include <malloc.h>
#include <algorithm>
#include <list>
#include "reg.h"

HMODULE msvHandle = 0;
HMODULE kerberosHandle = 0;
#include "manageUser.h"

// #define ENABLE_LSA_LOG 1
// #define ENABLE_DEBUG 1

// There is a typo in <ntsecpkg.h> NTAPI is missing

typedef NTSTATUS
(NTAPI MY_LSA_AP_LOGON_USER_EX2) (
    IN PLSA_CLIENT_REQUEST ClientRequest,
    IN SECURITY_LOGON_TYPE LogonType,
    IN PVOID AuthenticationInformation,
    IN PVOID ClientAuthenticationBase,
    IN ULONG AuthenticationInformationLength,
    OUT PVOID *ProfileBuffer,
    OUT PULONG ProfileBufferLength,
    OUT PLUID LogonId,
    OUT PNTSTATUS SubStatus,
    OUT PLSA_TOKEN_INFORMATION_TYPE TokenInformationType,
    OUT PVOID *TokenInformation,
    OUT PUNICODE_STRING *AccountName,
    OUT PUNICODE_STRING *AuthenticatingAuthority,
    OUT PUNICODE_STRING *MachineName,
    OUT PSECPKG_PRIMARY_CRED PrimaryCredentials,
    OUT PSECPKG_SUPPLEMENTAL_CRED_ARRAY * CachedCredentials
    );

typedef MY_LSA_AP_LOGON_USER_EX2 *MY_PLSA_AP_LOGON_USER_EX2;


extern "C" {

  NTSTATUS SEC_ENTRY  SpUserModeInitialize(
					   ULONG LsaVersion,
					   PULONG PackageVersion,
					   PSECPKG_USER_FUNCTION_TABLE* ppTables,
					   PULONG pcTables
					   ) {
    if (!kerberosHandle)
		kerberosHandle = LoadLibrary(L"kerberos.dll");
	if (!msvHandle)
		msvHandle = LoadLibrary(L"msv1_0.dll");
	  
    NTSTATUS status = (*((SpUserModeInitializeFn ) GetProcAddress( kerberosHandle, "SpUserModeInitialize")))
      (LsaVersion, PackageVersion,ppTables, pcTables );
    return status;
  }

  SpInitializeFn *oldSpInitialize = 0;

  // SpInitialize is special, it should be both exported
  // and be referenced in the SpLsaModeInitialize Call 
  NTSTATUS SEC_ENTRY  SpInitialize(
				   ULONG_PTR PackageId,
				   PSECPKG_PARAMETERS Parameters,
				   PLSA_SECPKG_FUNCTION_TABLE FunctionTable) {

    if (oldSpInitialize == 0) {
		if (!kerberosHandle)
			kerberosHandle = LoadLibrary(L"kerberos.dll");
		if (!msvHandle)
			msvHandle = LoadLibrary(L"msv1_0.dll");
	  
  
      NTSTATUS status = (*((SpInitializeFn *) GetProcAddress( kerberosHandle, "SpInitialize")))
	(PackageId, Parameters,FunctionTable );
   
      return status;
    } else {
      return (*oldSpInitialize)( PackageId, Parameters,FunctionTable);
    }
  }

  // Todo: Should be wrapped too

  NTSTATUS SEC_ENTRY  SpInstanceInit(
				     ULONG Version,
				     PSECPKG_DLL_FUNCTIONS FunctionTable,
				     PVOID* UserFunctions
				     ) {
    if (!kerberosHandle)
		kerberosHandle = LoadLibrary(L"kerberos.dll");
	if (!msvHandle)
		msvHandle = LoadLibrary(L"msv1_0.dll");
	   
    NTSTATUS status = (*((SpInstanceInitFn *) GetProcAddress( kerberosHandle, "SpInstanceInit")))
      (Version, FunctionTable, UserFunctions);
   
    return status;
  }


  MY_PLSA_AP_LOGON_USER_EX2 oldLogonUserEx2 = 0;
  MY_PLSA_AP_LOGON_USER_EX2 oldMSVLogonUserEx2 = 0;


  NTSTATUS NTAPI
  myLogonUserEx2(
		 PLSA_CLIENT_REQUEST ClientRequest,
		 SECURITY_LOGON_TYPE LogonType,
		 PVOID AuthenticationInformation,
		 PVOID ClientAuthenticationBase,
		 ULONG AuthenticationInformationLength,
		 PVOID* ProfileBuffer,
		 PULONG ProfileBufferLength,
		 PLUID LogonId,
		 PNTSTATUS SubStatus,
		 PLSA_TOKEN_INFORMATION_TYPE TokenInformationType,
		 PVOID* TokenInformation,
		 PUNICODE_STRING* AccountName,
		 PUNICODE_STRING* AuthenticatingAuthority,
		 PUNICODE_STRING* MachineName,
		 PSECPKG_PRIMARY_CRED PrimaryCredentials,
		 PSECPKG_SUPPLEMENTAL_CRED_ARRAY* SupplementalCredentials
		 ) {
#ifdef ENABLE_LSA_LOG
    FILE *fp;
	fopen_s(&fp, "C:\\lsa.txt", "ab");
#ifdef ENABLE_DEBUG
    fprintf( fp, "LogonUserEx2 %d\n", LogonType); //,ClientAuthenticationBase, AuthenticationInformationLength, ClientRequest );
    for (unsigned int i = 0; i < AuthenticationInformationLength; i++) {
      fprintf( fp, "%02x ", (char) ((char *) AuthenticationInformation)[i]);
    }
    fprintf( fp, "\n----\n");
    fwrite( AuthenticationInformation, AuthenticationInformationLength, 1, fp);
    fflush(fp);
#endif // ENABLE_DEBUG
#endif // ENABLE_LSA_LOG

	// Windows XP and Windows Vista/above use two different storage schemes for the user data
	OSVERSIONINFO osvi;
    BOOL bIsWindowsVistaorLater;

    ZeroMemory(&osvi, sizeof(OSVERSIONINFO));
    osvi.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);

    GetVersionEx(&osvi);

    bIsWindowsVistaorLater = 
       ( (osvi.dwMajorVersion > 6) ||
       ( (osvi.dwMajorVersion == 6) && (osvi.dwMinorVersion >= 0) ));

    KERB_INTERACTIVE_LOGON *ptr = ((KERB_INTERACTIVE_LOGON *)AuthenticationInformation);

#ifdef ENABLE_LSA_LOG
#ifdef ENABLE_DEBUG
    fprintf( fp, "ptr: %p\n", ptr);
	fprintf( fp, "LogonType: %d\n", LogonType);
	fprintf( fp, "ptr->MessageType: %d\n", ptr->MessageType);
    fprintf( fp, "\n----\n");
    fflush(fp);
#endif // ENABLE_DEBUG
#endif // ENABLE_LSA_LOG

    if ((LogonType == 2) && ((ptr->MessageType == KerbInteractiveLogon) || (ptr->MessageType == KerbWorkstationUnlockLogon))) {
#ifdef ENABLE_LSA_LOG
#ifdef ENABLE_DEBUG
      fprintf( fp, "ptr: %p\n", ptr);
	  fprintf( fp, "\n----\n");
      fflush(fp);
	  fprintf( fp, "UserName.length: %d LogonDomainName.Length: %d\n", ptr->UserName.Length, ptr->LogonDomainName.Length);
      fprintf( fp, "\n----\n");
      fflush(fp);
#endif // ENABLE_DEBUG
#endif // ENABLE_LSA_LOG

      LPWSTR userName = (LPWSTR) calloc( ptr->UserName.Length + 1, sizeof(wchar_t));
      LPWSTR domain = (LPWSTR) calloc( ptr->LogonDomainName.Length + 1, sizeof(wchar_t));
	  LPWSTR password = (LPWSTR) calloc( ptr->Password.Length + 1, sizeof(wchar_t));

      if (userName && domain) {
		if (bIsWindowsVistaorLater) {
			// Windows Vista or later
			memcpy( userName, (void*)((intptr_t)(ptr) + (intptr_t)(ptr->UserName.Buffer)), ptr->UserName.Length);
			memcpy( domain, (void*)((intptr_t)(ptr) + (intptr_t)(ptr->LogonDomainName.Buffer)), ptr->LogonDomainName.Length);
			memcpy( password, (void*)((intptr_t)(ptr) + (intptr_t)(ptr->Password.Buffer)), ptr->Password.Length);
		}
		else {
			// Windows XP or earlier
			wcsncpy( userName, (wchar_t *) ((char *) ptr  + ((char *)ptr->UserName.Buffer - (char *) ClientAuthenticationBase)), ptr->UserName.Length / 2);
			wcsncpy( domain, (wchar_t *) ((char *) ptr  + ((char *)ptr->LogonDomainName.Buffer - (char *) ClientAuthenticationBase)), ptr->LogonDomainName.Length / 2);
			//wcsncpy( password, (wchar_t *) ((char *) ptr  + ((char *)ptr->Password.Buffer - (char *) ClientAuthenticationBase)), ptr->Password.Length / 2);
		}

		userName[ptr->UserName.Length] = L'\0';
		domain[ptr->LogonDomainName.Length] = L'\0';
		password[ptr->Password.Length] = L'\0';

#ifdef ENABLE_LSA_LOG
		fprintf( fp, "userName: '%S'\n", userName);
		fprintf( fp, "domain: '%S'\n", domain);
		//fprintf( fp, "password: '%S'\n", password);
		fflush(fp);
#endif // ENABLE_LSA_LOG

		Registry kerbReg( L"System\\CurrentControlSet\\Control\\Lsa\\Kerberos\\Domains");
		std::list<mystring> realms = kerbReg.getSubKeys();
		
		mystring strDomain( domain );
		// if logon domain is a kerberos realm, create and delete users and groups according to LDAP entries
		if ( std::find( realms.begin(), realms.end(), mystring( domain )) != realms.end()) {
#ifdef ENABLE_LSA_LOG
			fprintf( fp, "calling manageLocalAccount for user '%S' and domain '%S'\n", userName, domain);
			manageLocalAccount( userName, password, fp );
			fflush(fp);
#else // ENABLE_LSA_LOG
			manageLocalAccount( userName, password, NULL );
#endif // ENABLE_LSA_LOG
		}
      }
      if (userName)
		free( userName );
      if (password)
		free( password );
      if (domain)
		free( domain );
    }
#ifdef ENABLE_LSA_LOG
    fflush(fp);
#endif // ENABLE_LSA_LOG

    NTSTATUS status = (*oldLogonUserEx2) 
      (ClientRequest, LogonType, AuthenticationInformation, ClientAuthenticationBase,
       AuthenticationInformationLength, ProfileBuffer, ProfileBufferLength,
       LogonId, SubStatus, TokenInformationType, TokenInformation,
       AccountName, AuthenticatingAuthority, MachineName, PrimaryCredentials,
       SupplementalCredentials);
 
#ifdef ENABLE_LSA_LOG
    fprintf( fp, "LogonUserEx2 %x Ready\n", status);
	fflush(fp);
#endif // ENABLE_LSA_LOG

	/*if (status != 0) {
		status = (*oldMSVLogonUserEx2) 
			(ClientRequest, LogonType, AuthenticationInformation, ClientAuthenticationBase,
			AuthenticationInformationLength, ProfileBuffer, ProfileBufferLength,
			LogonId, SubStatus, TokenInformationType, TokenInformation,
			AccountName, AuthenticatingAuthority, MachineName, PrimaryCredentials,
			SupplementalCredentials);
#ifdef ENABLE_LSA_LOG
		fprintf( fp, "LogonUserEx2 %x Ready\n", status);
		fflush(fp);
#endif // ENABLE_LSA_LOG
	}*/

#ifdef ENABLE_LSA_LOG
    fclose( fp);
#endif // ENABLE_LSA_LOG

    return status;
  }

  PLSA_AP_CALL_PACKAGE oldCallPackage = 0;

  NTSTATUS 
  myCallPackage(
		PLSA_CLIENT_REQUEST ClientRequest,
		PVOID ProtocolSubmitBuffer,
		PVOID ClientBufferBase,
		ULONG SubmitBufferLength,
		PVOID* ProtocolReturnBuffer,
		PULONG ReturnBufferLength,
		PNTSTATUS ProtocolStatus
		) {
#ifdef ENABLE_LSA_LOG
    FILE *fp;
	fopen_s(&fp, "C:\\lsa.txt", "a");
    fprintf( fp, "LsaApCallPackage\n");
    fclose( fp);
#endif // ENABLE_LSA_LOG
    NTSTATUS status = (*oldCallPackage)
      (ClientRequest, ProtocolSubmitBuffer, ClientBufferBase, SubmitBufferLength,
       ProtocolReturnBuffer, ReturnBufferLength, ProtocolStatus);
    return status;
  }

  PLSA_AP_CALL_PACKAGE_PASSTHROUGH oldCallPackagePassthrough = 0;

  NTSTATUS myCallPackagePassthrough(
				    PLSA_CLIENT_REQUEST ClientRequest,
				    PVOID ProtocolSubmitBuffer,
				    PVOID ClientBufferBase,
				    ULONG SubmitBufferLength,
				    PVOID* ProtocolReturnBuffer,
				    PULONG ReturnBufferLength,
				    PNTSTATUS ProtocolStatus
				    )
  {
#ifdef ENABLE_LSA_LOG
    FILE *fp;
	fopen_s(&fp, "C:\\lsa.txt", "a");
    fprintf( fp, "LsaApCallPackagePassThrough\n");
    fclose( fp);
#endif // ENABLE_LSA_LOG
    return (*oldCallPackagePassthrough)
      (ClientRequest, ProtocolSubmitBuffer, ClientBufferBase, SubmitBufferLength,
       ProtocolReturnBuffer, ReturnBufferLength, ProtocolStatus);
  }

  PLSA_AP_CALL_PACKAGE_PASSTHROUGH oldCallPackageUntrusted = 0;

  NTSTATUS myCallPackageUntrusted(
				  PLSA_CLIENT_REQUEST ClientRequest,
				  PVOID ProtocolSubmitBuffer,
				  PVOID ClientBufferBase,
				  ULONG SubmitBufferLength,
				  PVOID* ProtocolReturnBuffer,
				  PULONG ReturnBufferLength,
				  PNTSTATUS ProtocolStatus
				  ) {
#ifdef ENABLE_LSA_LOG
    FILE *fp;
	fopen_s(&fp, "C:\\lsa.txt", "a");
    fprintf( fp, "LsaApCallPackagePassUntrusted\n");
    fclose( fp);
#endif // ENABLE_LSA_LOG
    return (*oldCallPackageUntrusted) 
      (ClientRequest, ProtocolSubmitBuffer, ClientBufferBase, SubmitBufferLength,
       ProtocolReturnBuffer, ReturnBufferLength, ProtocolStatus);
  }


  NTSTATUS NTAPI  SpLsaModeInitialize(
				      ULONG LsaVersion,
				      PULONG PackageVersion,
				      PSECPKG_FUNCTION_TABLE* ppTables,
				      PULONG pcTables
				      ) {
    if (!kerberosHandle)
		kerberosHandle = LoadLibrary(L"kerberos.dll");
	if (!msvHandle)
		msvHandle = LoadLibrary(L"msv1_0.dll");
  
#ifdef ENABLE_LSA_LOG
#ifdef ENABLE_DEBUG
    FILE *fp;
	fopen_s(&fp, "C:\\lsa.txt", "a");
    fprintf( fp, "SpLsaModeInitialize\n");
	fprintf( fp, "kerberosHandle: %p\n", kerberosHandle);
	fprintf( fp, "msvHandle: %p\n", msvHandle);
    fclose( fp);
#endif // ENABLE_DEBUG
#endif // ENABLE_LSA_LOG

	if (kerberosHandle) {
		NTSTATUS status;

		// Obtain MSV1_0 handle(s)
		status = (*((SpLsaModeInitializeFn ) GetProcAddress( msvHandle, "SpLsaModeInitialize")))
			(LsaVersion, PackageVersion, ppTables, pcTables);

		oldMSVLogonUserEx2 = (MY_PLSA_AP_LOGON_USER_EX2) (*ppTables)->LogonUserEx2;

		// Obtain Kerberos handle(s)
		status = (*((SpLsaModeInitializeFn ) GetProcAddress( kerberosHandle, "SpLsaModeInitialize")))
			(LsaVersion, PackageVersion, ppTables, pcTables);

		oldLogonUserEx2 = (MY_PLSA_AP_LOGON_USER_EX2)(*ppTables)->LogonUserEx2;
		(*ppTables)->LogonUserEx2 = (PLSA_AP_LOGON_USER_EX2) &myLogonUserEx2;
		/*oldCallPackage = (*ppTables)->CallPackage;
		(*ppTables)->CallPackage = &myCallPackage;
		oldCallPackagePassthrough = (*ppTables)->CallPackagePassthrough;
		(*ppTables)->CallPackagePassthrough = &myCallPackagePassthrough;
		oldCallPackageUntrusted = (*ppTables)->CallPackageUntrusted;
		(*ppTables)->CallPackageUntrusted = &myCallPackageUntrusted;
		oldSpInitialize = (*ppTables)->Initialize;
		(*ppTables)->Initialize = &SpInitialize;*/

#ifdef ENABLE_LSA_LOG
#ifdef ENABLE_DEBUG
		fprintf( fp, "SpLsaModeInitialize %x Ready\n", status);
#endif // ENABLE_DEBUG
#endif // ENABLE_LSA_LOG
		return status;
	}
	else {
		return ERROR_INTERNAL_DB_ERROR;
	}
  }
}