# SECUREAUTH LABS. Copyright 2018 SecureAuth Corporation. All rights reserved.
# Copyright (c) 2017 @MrAnde7son
#
# This software is provided under under a slightly modified version
# of the Apache Software License. See the accompanying LICENSE file
# for more information.
#
# Author: Itamar (@MrAnde7son)
#
# Description:
#   Initial [MS-EVEN6] Interface implementation
#
#   Best way to learn how to use these calls is to grab the protocol standard
#   so you understand what the call does, and then read the test case located
#   at https://github.com/SecureAuthCorp/impacket/tree/master/tests/SMB_RPC
#
#   Some calls have helper functions, which makes it even easier to use.
#   They are located at the end of this file.
#   Helper functions start with "h"<name of the call>.
#   There are test cases for them too.
#
from aiosmb.dcerpc.v5 import system_errors
from aiosmb.dcerpc.v5.dtypes import WSTR, DWORD, LPWSTR, ULONG, LARGE_INTEGER, WORD, BYTE
from aiosmb.dcerpc.v5.ndr import NDRCALL, NDRPOINTER, NDRUniConformantArray, NDRUniVaryingArray, NDRSTRUCT
from aiosmb.dcerpc.v5.rpcrt import DCERPCException
from aiosmb.dcerpc.v5.uuid import uuidtup_to_bin

MSRPC_UUID_EVEN6 = uuidtup_to_bin(('F6BEAFF7-1E19-4FBB-9F8F-B89E2018337C', '1.0'))

class DCERPCSessionError(DCERPCException):
	def __init__(self, error_string=None, error_code=None, packet=None):
		DCERPCException.__init__(self, error_string, error_code, packet)


################################################################################
# CONSTANTS
################################################################################

# Evt Path Flags
EvtQueryChannelName = 0x00000001
EvtQueryFilePath = 0x00000002
EvtReadOldestToNewest = 0x00000100
EvtReadNewestToOldest = 0x00000200

################################################################################
# STRUCTURES
################################################################################

class CONTEXT_HANDLE_LOG_HANDLE(NDRSTRUCT):
	align = 1
	structure = (
		('Data', '20s=""'),
	)

class PCONTEXT_HANDLE_LOG_HANDLE(NDRPOINTER):
	referent = (
		('Data', CONTEXT_HANDLE_LOG_HANDLE),
	)

class CONTEXT_HANDLE_LOG_QUERY(NDRSTRUCT):
	align = 1
	structure = (
		('Data', '20s=""'),
	)

class PCONTEXT_HANDLE_LOG_QUERY(NDRPOINTER):
	referent = (
		('Data', CONTEXT_HANDLE_LOG_QUERY),
	)

class LPPCONTEXT_HANDLE_LOG_QUERY(NDRPOINTER):
	referent = (
		('Data', PCONTEXT_HANDLE_LOG_QUERY),
	)

class CONTEXT_HANDLE_OPERATION_CONTROL(NDRSTRUCT):
	align = 1
	structure = (
		('Data', '20s=""'),
	)

class PCONTEXT_HANDLE_OPERATION_CONTROL(NDRPOINTER):
	referent = (
		('Data', CONTEXT_HANDLE_OPERATION_CONTROL),
	)

# 2.2.11 EvtRpcQueryChannelInfo
class EvtRpcQueryChannelInfo(NDRSTRUCT):
	structure = (
		('Name', LPWSTR),
		('Status', DWORD),
	)

class EvtRpcQueryChannelInfoArray(NDRUniVaryingArray):
	item = EvtRpcQueryChannelInfo

class LPEvtRpcQueryChannelInfoArray(NDRPOINTER):
	referent = (
		('Data', EvtRpcQueryChannelInfoArray)
	)

class RPC_INFO(NDRSTRUCT):
	structure = (
		('Error', DWORD),
		('SubError', DWORD),
		('SubErrorParam', DWORD),
	)

class PRPC_INFO(NDRPOINTER):
	referent = (
		('Data', RPC_INFO)
	)

class LPWSTR_ARRAY(NDRUniConformantArray):
	item = LPWSTR

class PLPWSTR_ARRAY(NDRPOINTER):
	referent = (
		('Data',LPWSTR_ARRAY),
	)

class WSTR_ARRAY(NDRUniVaryingArray):
	item = WSTR

class DWORD_ARRAY(NDRUniVaryingArray):
	item = DWORD

class LPDWORD_ARRAY(NDRPOINTER):
	referent = (
		('Data', DWORD_ARRAY)
	)

class BYTE_ARRAY(NDRUniVaryingArray):
	item = 'c'

class CBYTE_ARRAY(NDRUniVaryingArray):
	item = BYTE

class CDWORD_ARRAY(NDRUniConformantArray):
	item = DWORD

class LPBYTE_ARRAY(NDRPOINTER):
	referent = (
		('Data', CBYTE_ARRAY)
	)

class ULONG_ARRAY(NDRUniVaryingArray):
	item = ULONG

# 2.3.1 EVENT_DESCRIPTOR
class EVENT_DESCRIPTOR(NDRSTRUCT):
	structure = (
		('Id', WORD),
		('Version', BYTE),
		('Channel', BYTE),
		('LevelSeverity', BYTE),
		('Opcode', BYTE),
		('Task', WORD),
		('Keyword', ULONG),
	)

class BOOKMARK(NDRSTRUCT):
	structure = (
		('BookmarkSize', DWORD),
		('HeaderSize', '<L=0x18'),
		('ChannelSize', DWORD),
		('CurrentChannel', DWORD),
		('ReadDirection', DWORD),
		('RecordIdsOffset', DWORD),
		('LogRecordNumbers', ULONG_ARRAY),
	)


#2.2.17 RESULT_SET
class RESULT_SET(NDRSTRUCT):
	structure = (
		('TotalSize', DWORD),
		('HeaderSize', DWORD),
		('EventOffset', DWORD),
		('BookmarkOffset', DWORD),
		('BinXmlSize', DWORD),
		('EventData', BYTE_ARRAY),
		#('NumberOfSubqueryIDs', '<L=0'),
		#('SubqueryIDs', BYTE_ARRAY),
		#('BookMarkData', BOOKMARK),
	)

################################################################################
# RPC CALLS
################################################################################

class EvtRpcRegisterLogQuery(NDRCALL):
	opnum = 5
	structure = (
		('Path', LPWSTR),
		('Query', WSTR),
		('Flags', DWORD),
	)

class EvtRpcRegisterLogQueryResponse(NDRCALL):
	structure = (
		('Handle', CONTEXT_HANDLE_LOG_QUERY),
		('OpControl', CONTEXT_HANDLE_OPERATION_CONTROL),
		('QueryChannelInfoSize', DWORD),
		('QueryChannelInfo', EvtRpcQueryChannelInfoArray),
		('Error', RPC_INFO),
		)

class EvtRpcQueryNext(NDRCALL):
	opnum = 11
	structure = (
		('LogQuery', CONTEXT_HANDLE_LOG_QUERY),
		('NumRequestedRecords', DWORD),
		('TimeOutEnd', DWORD),
		('Flags', DWORD),
	)

class EvtRpcQueryNextResponse(NDRCALL):
	structure = (
		('NumActualRecords', DWORD),
		('EventDataIndices', DWORD_ARRAY),
		('EventDataSizes', DWORD_ARRAY),
		('ResultBufferSize', DWORD),
		('ResultBuffer', BYTE_ARRAY),
		('ErrorCode', ULONG),
	)

class EvtRpcQuerySeek(NDRCALL):
	opnum = 12
	structure = (
		('LogQuery', CONTEXT_HANDLE_LOG_QUERY),
		('Pos', LARGE_INTEGER),
		('BookmarkXML', LPWSTR),
		('Flags', DWORD),
	)

class EvtRpcQuerySeekResponse(NDRCALL):
	structure = (
		('Error', RPC_INFO),
	)

class EvtRpcClose(NDRCALL):
	opnum = 13
	structure = (
		("Handle", CONTEXT_HANDLE_LOG_HANDLE),
	)

class EvtRpcCloseResponse(NDRCALL):
	structure = (
		("Handle", PCONTEXT_HANDLE_LOG_HANDLE),
		('ErrorCode', ULONG),
	)

class EvtRpcOpenLogHandle(NDRCALL):
	opnum = 17
	structure = (
		('Channel', WSTR),
		('Flags', DWORD),
	)

class EvtRpcOpenLogHandleResponse(NDRCALL):
	structure = (
		('Handle', PCONTEXT_HANDLE_LOG_HANDLE),
		('Error', RPC_INFO),
	)

class EvtRpcGetChannelList(NDRCALL):
	opnum = 19
	structure = (
		('Flags', DWORD),
	)

class EvtRpcGetChannelListResponse(NDRCALL):
	structure = (
		('NumChannelPaths', DWORD),
		('ChannelPaths', PLPWSTR_ARRAY),
		('ErrorCode', ULONG),
	)

################################################################################
# OPNUMs and their corresponding structures
################################################################################

OPNUMS = {
	5   : (EvtRpcRegisterLogQuery, EvtRpcRegisterLogQueryResponse),
	11  : (EvtRpcQueryNext,  EvtRpcQueryNextResponse),
	12  : (EvtRpcQuerySeek, EvtRpcQuerySeekResponse),
	13  : (EvtRpcClose, EvtRpcCloseResponse),
	17  : (EvtRpcOpenLogHandle, EvtRpcOpenLogHandle),
	19  : (EvtRpcGetChannelList, EvtRpcGetChannelListResponse),
}

################################################################################
# HELPER FUNCTIONS
################################################################################

async def hEvtRpcRegisterLogQuery(dce, path, flags, query='*\x00'):
	request = EvtRpcRegisterLogQuery()

	request['Path'] = path
	request['Query'] = query
	request['Flags'] = flags
	resp = dce.request(request)
	return await resp

async def hEvtRpcQueryNext(dce, handle, numRequestedRecords, timeOutEnd=1000):
	request = EvtRpcQueryNext()

	request['LogQuery'] = handle
	request['NumRequestedRecords'] = numRequestedRecords
	request['TimeOutEnd'] = timeOutEnd
	request['Flags'] = 0
	return  await dce.request(request)

	#status = system_errors.ERROR_MORE_DATA
	#resp, err = await dce.request(request)
	#if err is not None:
	#	try:
	#		resp = await dce.request(request)
	#	except DCERPCException as e:
	#		if str(e).find('ERROR_NO_MORE_ITEMS') < 0:
	#			raise
	#		elif str(e).find('ERROR_TIMEOUT') < 0:
	#			raise
	#		resp = e.get_packet()
	#	return resp

async def hEvtRpcClose(dce, handle):
	request = EvtRpcClose()
	request['Handle'] = handle
	resp = dce.request(request)
	return await resp

async def hEvtRpcOpenLogHandle(dce, channel, flags):
	request = EvtRpcOpenLogHandle()

	request['Channel'] = channel
	request['Flags'] = flags
	return await dce.request(request)

async def hEvtRpcGetChannelList(dce):
	request = EvtRpcGetChannelList()
	request['Flags'] = 0
	return await dce.request(request)