1# Copyright (c) 2003-2016 CORE Security Technologies
2#
3# This software is provided under under a slightly modified version
4# of the Apache Software License. See the accompanying LICENSE file
5# for more information.
6#
7# Author: Alberto Solino (@agsolino)
8#
9# Description:
10#   [MS-RRP] Interface implementation
11#
12#   Best way to learn how to use these calls is to grab the protocol standard
13#   so you understand what the call does, and then read the test case located
14#   at https://github.com/CoreSecurity/impacket/tree/master/impacket/testcases/SMB_RPC
15#
16#   Some calls have helper functions, which makes it even easier to use.
17#   They are located at the end of this file.
18#   Helper functions start with "h"<name of the call>.
19#   There are test cases for them too.
20#
21from struct import unpack, pack
22
23from impacket.dcerpc.v5.ndr import NDRCALL, NDRSTRUCT, NDRPOINTER, NDRUniConformantVaryingArray, NDRUniConformantArray
24from impacket.dcerpc.v5.dtypes import DWORD, UUID, ULONG, LPULONG, BOOLEAN, SECURITY_INFORMATION, PFILETIME, \
25    RPC_UNICODE_STRING, FILETIME, NULL, MAXIMUM_ALLOWED, OWNER_SECURITY_INFORMATION, PWCHAR, PRPC_UNICODE_STRING
26from impacket.dcerpc.v5.rpcrt import DCERPCException
27from impacket import system_errors, LOG
28from impacket.uuid import uuidtup_to_bin
29
30MSRPC_UUID_RRP = uuidtup_to_bin(('338CD001-2244-31F1-AAAA-900038001003', '1.0'))
31
32class DCERPCSessionError(DCERPCException):
33    def __init__(self, error_string=None, error_code=None, packet=None):
34        DCERPCException.__init__(self, error_string, error_code, packet)
35
36    def __str__( self ):
37        key = self.error_code
38        if system_errors.ERROR_MESSAGES.has_key(key):
39            error_msg_short = system_errors.ERROR_MESSAGES[key][0]
40            error_msg_verbose = system_errors.ERROR_MESSAGES[key][1]
41            return 'RRP SessionError: code: 0x%x - %s - %s' % (self.error_code, error_msg_short, error_msg_verbose)
42        else:
43            return 'RRP SessionError: unknown error code: 0x%x' % self.error_code
44
45################################################################################
46# CONSTANTS
47################################################################################
48# 2.2.2 PREGISTRY_SERVER_NAME
49PREGISTRY_SERVER_NAME = PWCHAR
50
51# 2.2.3 error_status_t
52error_status_t = ULONG
53
54# 2.2.5 RRP_UNICODE_STRING
55RRP_UNICODE_STRING = RPC_UNICODE_STRING
56PRRP_UNICODE_STRING = PRPC_UNICODE_STRING
57
58# 2.2.4 REGSAM
59REGSAM = ULONG
60
61KEY_QUERY_VALUE        = 0x00000001
62KEY_SET_VALUE          = 0x00000002
63KEY_CREATE_SUB_KEY     = 0x00000004
64KEY_ENUMERATE_SUB_KEYS = 0x00000008
65KEY_CREATE_LINK        = 0x00000020
66KEY_WOW64_64KEY        = 0x00000100
67KEY_WOW64_32KEY        = 0x00000200
68
69REG_BINARY              = 3
70REG_DWORD               = 4
71REG_DWORD_LITTLE_ENDIAN = 4
72REG_DWORD_BIG_ENDIAN    = 5
73REG_EXPAND_SZ           = 2
74REG_LINK                = 6
75REG_MULTI_SZ            = 7
76REG_NONE                = 0
77REG_QWORD               = 11
78REG_QWORD_LITTLE_ENDIAN = 11
79REG_SZ                  = 1
80
81# 3.1.5.7 BaseRegCreateKey (Opnum 6)
82REG_CREATED_NEW_KEY     = 0x00000001
83REG_OPENED_EXISTING_KEY = 0x00000002
84
85# 3.1.5.19 BaseRegRestoreKey (Opnum 19)
86# Flags
87REG_WHOLE_HIVE_VOLATILE = 0x00000001
88REG_REFRESH_HIVE        = 0x00000002
89REG_NO_LAZY_FLUSH       = 0x00000004
90REG_FORCE_RESTORE       = 0x00000008
91
92################################################################################
93# STRUCTURES
94################################################################################
95# 2.2.1 RPC_HKEY
96class RPC_HKEY(NDRSTRUCT):
97    structure =  (
98        ('context_handle_attributes',ULONG),
99        ('context_handle_uuid',UUID),
100    )
101    def __init__(self, data = None,isNDR64 = False):
102        NDRSTRUCT.__init__(self, data, isNDR64)
103        self['context_handle_uuid'] = '\x00'*20
104
105# 2.2.6 RVALENT
106class RVALENT(NDRSTRUCT):
107    structure =  (
108        ('ve_valuename',PRRP_UNICODE_STRING),
109        ('ve_valuelen',DWORD),
110        ('ve_valueptr',DWORD),
111        ('ve_type',DWORD),
112    )
113
114class RVALENT_ARRAY(NDRUniConformantVaryingArray):
115    item = RVALENT
116
117# 2.2.9 RPC_SECURITY_DESCRIPTOR
118class BYTE_ARRAY(NDRUniConformantVaryingArray):
119    pass
120
121class PBYTE_ARRAY(NDRPOINTER):
122    referent = (
123        ('Data', BYTE_ARRAY),
124    )
125
126class RPC_SECURITY_DESCRIPTOR(NDRSTRUCT):
127    structure =  (
128        ('lpSecurityDescriptor',PBYTE_ARRAY),
129        ('cbInSecurityDescriptor',DWORD),
130        ('cbOutSecurityDescriptor',DWORD),
131    )
132
133# 2.2.8 RPC_SECURITY_ATTRIBUTES
134class RPC_SECURITY_ATTRIBUTES(NDRSTRUCT):
135    structure =  (
136        ('nLength',DWORD),
137        ('RpcSecurityDescriptor',RPC_SECURITY_DESCRIPTOR),
138        ('bInheritHandle',BOOLEAN),
139    )
140
141class PRPC_SECURITY_ATTRIBUTES(NDRPOINTER):
142    referent = (
143        ('Data', RPC_SECURITY_ATTRIBUTES),
144    )
145
146################################################################################
147# RPC CALLS
148################################################################################
149# 3.1.5.1 OpenClassesRoot (Opnum 0)
150class OpenClassesRoot(NDRCALL):
151    opnum = 0
152    structure = (
153       ('ServerName', PREGISTRY_SERVER_NAME),
154       ('samDesired', REGSAM),
155    )
156
157class OpenClassesRootResponse(NDRCALL):
158    structure = (
159       ('phKey', RPC_HKEY),
160       ('ErrorCode', error_status_t),
161    )
162
163# 3.1.5.2 OpenCurrentUser (Opnum 1)
164class OpenCurrentUser(NDRCALL):
165    opnum = 1
166    structure = (
167       ('ServerName', PREGISTRY_SERVER_NAME),
168       ('samDesired', REGSAM),
169    )
170
171class OpenCurrentUserResponse(NDRCALL):
172    structure = (
173       ('phKey', RPC_HKEY),
174       ('ErrorCode', error_status_t),
175    )
176
177# 3.1.5.3 OpenLocalMachine (Opnum 2)
178class OpenLocalMachine(NDRCALL):
179    opnum = 2
180    structure = (
181       ('ServerName', PREGISTRY_SERVER_NAME),
182       ('samDesired', REGSAM),
183    )
184
185class OpenLocalMachineResponse(NDRCALL):
186    structure = (
187       ('phKey', RPC_HKEY),
188       ('ErrorCode', error_status_t),
189    )
190
191# 3.1.5.4 OpenPerformanceData (Opnum 3)
192class OpenPerformanceData(NDRCALL):
193    opnum = 3
194    structure = (
195       ('ServerName', PREGISTRY_SERVER_NAME),
196       ('samDesired', REGSAM),
197    )
198
199class OpenPerformanceDataResponse(NDRCALL):
200    structure = (
201       ('phKey', RPC_HKEY),
202       ('ErrorCode', error_status_t),
203    )
204
205# 3.1.5.5 OpenUsers (Opnum 4)
206class OpenUsers(NDRCALL):
207    opnum = 4
208    structure = (
209       ('ServerName', PREGISTRY_SERVER_NAME),
210       ('samDesired', REGSAM),
211    )
212
213class OpenUsersResponse(NDRCALL):
214    structure = (
215       ('phKey', RPC_HKEY),
216       ('ErrorCode', error_status_t),
217    )
218
219# 3.1.5.6 BaseRegCloseKey (Opnum 5)
220class BaseRegCloseKey(NDRCALL):
221    opnum = 5
222    structure = (
223       ('hKey', RPC_HKEY),
224    )
225
226class BaseRegCloseKeyResponse(NDRCALL):
227    structure = (
228       ('hKey', RPC_HKEY),
229       ('ErrorCode', error_status_t),
230    )
231
232# 3.1.5.7 BaseRegCreateKey (Opnum 6)
233class BaseRegCreateKey(NDRCALL):
234    opnum = 6
235    structure = (
236       ('hKey', RPC_HKEY),
237       ('lpSubKey', RRP_UNICODE_STRING),
238       ('lpClass', RRP_UNICODE_STRING),
239       ('dwOptions', DWORD),
240       ('samDesired', REGSAM),
241       ('lpSecurityAttributes', PRPC_SECURITY_ATTRIBUTES),
242       ('lpdwDisposition', LPULONG),
243    )
244
245class BaseRegCreateKeyResponse(NDRCALL):
246    structure = (
247       ('phkResult', RPC_HKEY),
248       ('lpdwDisposition', LPULONG),
249       ('ErrorCode', error_status_t),
250    )
251
252# 3.1.5.8 BaseRegDeleteKey (Opnum 7)
253class BaseRegDeleteKey(NDRCALL):
254    opnum = 7
255    structure = (
256       ('hKey', RPC_HKEY),
257       ('lpSubKey', RRP_UNICODE_STRING),
258    )
259
260class BaseRegDeleteKeyResponse(NDRCALL):
261    structure = (
262       ('ErrorCode', error_status_t),
263    )
264
265# 3.1.5.9 BaseRegDeleteValue (Opnum 8)
266class BaseRegDeleteValue(NDRCALL):
267    opnum = 8
268    structure = (
269       ('hKey', RPC_HKEY),
270       ('lpValueName', RRP_UNICODE_STRING),
271    )
272
273class BaseRegDeleteValueResponse(NDRCALL):
274    structure = (
275       ('ErrorCode', error_status_t),
276    )
277
278# 3.1.5.10 BaseRegEnumKey (Opnum 9)
279class BaseRegEnumKey(NDRCALL):
280    opnum = 9
281    structure = (
282       ('hKey', RPC_HKEY),
283       ('dwIndex', DWORD),
284       ('lpNameIn', RRP_UNICODE_STRING),
285       ('lpClassIn', PRRP_UNICODE_STRING),
286       ('lpftLastWriteTime', PFILETIME),
287    )
288
289class BaseRegEnumKeyResponse(NDRCALL):
290    structure = (
291       ('lpNameOut', RRP_UNICODE_STRING),
292       ('lplpClassOut', PRRP_UNICODE_STRING),
293       ('lpftLastWriteTime', PFILETIME),
294       ('ErrorCode', error_status_t),
295    )
296
297# 3.1.5.11 BaseRegEnumValue (Opnum 10)
298class BaseRegEnumValue(NDRCALL):
299    opnum = 10
300    structure = (
301       ('hKey', RPC_HKEY),
302       ('dwIndex', DWORD),
303       ('lpValueNameIn', RRP_UNICODE_STRING),
304       ('lpType', LPULONG),
305       ('lpData', PBYTE_ARRAY),
306       ('lpcbData', LPULONG),
307       ('lpcbLen', LPULONG),
308    )
309
310class BaseRegEnumValueResponse(NDRCALL):
311    structure = (
312       ('lpValueNameOut', RRP_UNICODE_STRING),
313       ('lpType', LPULONG),
314       ('lpData', PBYTE_ARRAY),
315       ('lpcbData', LPULONG),
316       ('lpcbLen', LPULONG),
317       ('ErrorCode', error_status_t),
318    )
319
320# 3.1.5.12 BaseRegFlushKey (Opnum 11)
321class BaseRegFlushKey(NDRCALL):
322    opnum = 11
323    structure = (
324       ('hKey', RPC_HKEY),
325    )
326
327class BaseRegFlushKeyResponse(NDRCALL):
328    structure = (
329       ('ErrorCode', error_status_t),
330    )
331
332# 3.1.5.13 BaseRegGetKeySecurity (Opnum 12)
333class BaseRegGetKeySecurity(NDRCALL):
334    opnum = 12
335    structure = (
336       ('hKey', RPC_HKEY),
337       ('SecurityInformation', SECURITY_INFORMATION),
338       ('pRpcSecurityDescriptorIn', RPC_SECURITY_DESCRIPTOR),
339    )
340
341class BaseRegGetKeySecurityResponse(NDRCALL):
342    structure = (
343       ('pRpcSecurityDescriptorOut', RPC_SECURITY_DESCRIPTOR),
344       ('ErrorCode', error_status_t),
345    )
346
347# 3.1.5.14 BaseRegLoadKey (Opnum 13)
348class BaseRegLoadKey(NDRCALL):
349    opnum = 13
350    structure = (
351       ('hKey', RPC_HKEY),
352       ('lpSubKey', RRP_UNICODE_STRING),
353       ('lpFile', RRP_UNICODE_STRING),
354    )
355
356class BaseRegLoadKeyResponse(NDRCALL):
357    structure = (
358       ('ErrorCode', error_status_t),
359    )
360
361# 3.1.5.15 BaseRegOpenKey (Opnum 15)
362class BaseRegOpenKey(NDRCALL):
363    opnum = 15
364    structure = (
365       ('hKey', RPC_HKEY),
366       ('lpSubKey', RRP_UNICODE_STRING),
367       ('dwOptions', DWORD),
368       ('samDesired', REGSAM),
369    )
370
371class BaseRegOpenKeyResponse(NDRCALL):
372    structure = (
373       ('phkResult', RPC_HKEY),
374       ('ErrorCode', error_status_t),
375    )
376
377# 3.1.5.16 BaseRegQueryInfoKey (Opnum 16)
378class BaseRegQueryInfoKey(NDRCALL):
379    opnum = 16
380    structure = (
381       ('hKey', RPC_HKEY),
382       ('lpClassIn', RRP_UNICODE_STRING),
383    )
384
385class BaseRegQueryInfoKeyResponse(NDRCALL):
386    structure = (
387       ('lpClassOut', RPC_UNICODE_STRING),
388       ('lpcSubKeys', DWORD),
389       ('lpcbMaxSubKeyLen', DWORD),
390       ('lpcbMaxClassLen', DWORD),
391       ('lpcValues', DWORD),
392       ('lpcbMaxValueNameLen', DWORD),
393       ('lpcbMaxValueLen', DWORD),
394       ('lpcbSecurityDescriptor', DWORD),
395       ('lpftLastWriteTime', FILETIME),
396       ('ErrorCode', error_status_t),
397    )
398
399# 3.1.5.17 BaseRegQueryValue (Opnum 17)
400class BaseRegQueryValue(NDRCALL):
401    opnum = 17
402    structure = (
403       ('hKey', RPC_HKEY),
404       ('lpValueName', RRP_UNICODE_STRING),
405       ('lpType', LPULONG),
406       ('lpData', PBYTE_ARRAY),
407       ('lpcbData', LPULONG),
408       ('lpcbLen', LPULONG),
409    )
410
411class BaseRegQueryValueResponse(NDRCALL):
412    structure = (
413       ('lpType', LPULONG),
414       ('lpData', PBYTE_ARRAY),
415       ('lpcbData', LPULONG),
416       ('lpcbLen', LPULONG),
417       ('ErrorCode', error_status_t),
418    )
419
420# 3.1.5.18 BaseRegReplaceKey (Opnum 18)
421class BaseRegReplaceKey(NDRCALL):
422    opnum = 18
423    structure = (
424       ('hKey', RPC_HKEY),
425       ('lpSubKey', RRP_UNICODE_STRING),
426       ('lpNewFile', RRP_UNICODE_STRING),
427       ('lpOldFile', RRP_UNICODE_STRING),
428    )
429
430class BaseRegReplaceKeyResponse(NDRCALL):
431    structure = (
432       ('ErrorCode', error_status_t),
433    )
434
435# 3.1.5.19 BaseRegRestoreKey (Opnum 19)
436class BaseRegRestoreKey(NDRCALL):
437    opnum = 19
438    structure = (
439       ('hKey', RPC_HKEY),
440       ('lpFile', RRP_UNICODE_STRING),
441       ('Flags', DWORD),
442    )
443
444class BaseRegRestoreKeyResponse(NDRCALL):
445    structure = (
446       ('ErrorCode', error_status_t),
447    )
448
449# 3.1.5.20 BaseRegSaveKey (Opnum 20)
450class BaseRegSaveKey(NDRCALL):
451    opnum = 20
452    structure = (
453       ('hKey', RPC_HKEY),
454       ('lpFile', RRP_UNICODE_STRING),
455       ('pSecurityAttributes', PRPC_SECURITY_ATTRIBUTES),
456    )
457
458class BaseRegSaveKeyResponse(NDRCALL):
459    structure = (
460       ('ErrorCode', error_status_t),
461    )
462
463# 3.1.5.21 BaseRegSetKeySecurity (Opnum 21)
464class BaseRegSetKeySecurity(NDRCALL):
465    opnum = 21
466    structure = (
467       ('hKey', RPC_HKEY),
468       ('SecurityInformation', SECURITY_INFORMATION),
469       ('pRpcSecurityDescriptor', RPC_SECURITY_DESCRIPTOR),
470    )
471
472class BaseRegSetKeySecurityResponse(NDRCALL):
473    structure = (
474       ('ErrorCode', error_status_t),
475    )
476
477# 3.1.5.22 BaseRegSetValue (Opnum 22)
478class BaseRegSetValue(NDRCALL):
479    opnum = 22
480    structure = (
481       ('hKey', RPC_HKEY),
482       ('lpValueName', RRP_UNICODE_STRING),
483       ('dwType', DWORD),
484       ('lpData', NDRUniConformantArray),
485       ('cbData', DWORD),
486    )
487
488class BaseRegSetValueResponse(NDRCALL):
489    structure = (
490       ('ErrorCode', error_status_t),
491    )
492
493# 3.1.5.23 BaseRegUnLoadKey (Opnum 23)
494class BaseRegUnLoadKey(NDRCALL):
495    opnum = 23
496    structure = (
497       ('hKey', RPC_HKEY),
498       ('lpSubKey', RRP_UNICODE_STRING),
499    )
500
501class BaseRegUnLoadKeyResponse(NDRCALL):
502    structure = (
503       ('ErrorCode', error_status_t),
504    )
505
506# 3.1.5.24 BaseRegGetVersion (Opnum 26)
507class BaseRegGetVersion(NDRCALL):
508    opnum = 26
509    structure = (
510       ('hKey', RPC_HKEY),
511    )
512
513class BaseRegGetVersionResponse(NDRCALL):
514    structure = (
515       ('lpdwVersion', DWORD),
516       ('ErrorCode', error_status_t),
517    )
518
519# 3.1.5.25 OpenCurrentConfig (Opnum 27)
520class OpenCurrentConfig(NDRCALL):
521    opnum = 27
522    structure = (
523       ('ServerName', PREGISTRY_SERVER_NAME),
524       ('samDesired', REGSAM),
525    )
526
527class OpenCurrentConfigResponse(NDRCALL):
528    structure = (
529       ('phKey', RPC_HKEY),
530       ('ErrorCode', error_status_t),
531    )
532
533# 3.1.5.26 BaseRegQueryMultipleValues (Opnum 29)
534class BaseRegQueryMultipleValues(NDRCALL):
535    opnum = 29
536    structure = (
537       ('hKey', RPC_HKEY),
538       ('val_listIn', RVALENT_ARRAY),
539       ('num_vals', DWORD),
540       ('lpvalueBuf', PBYTE_ARRAY),
541       ('ldwTotsize', DWORD),
542    )
543
544class BaseRegQueryMultipleValuesResponse(NDRCALL):
545    structure = (
546       ('val_listOut', RVALENT_ARRAY),
547       ('lpvalueBuf', PBYTE_ARRAY),
548       ('ldwTotsize', DWORD),
549       ('ErrorCode', error_status_t),
550    )
551
552# 3.1.5.27 BaseRegSaveKeyEx (Opnum 31)
553class BaseRegSaveKeyEx(NDRCALL):
554    opnum = 31
555    structure = (
556       ('hKey', RPC_HKEY),
557       ('lpFile', RRP_UNICODE_STRING),
558       ('pSecurityAttributes', PRPC_SECURITY_ATTRIBUTES),
559       ('Flags', DWORD),
560    )
561
562class BaseRegSaveKeyExResponse(NDRCALL):
563    structure = (
564       ('ErrorCode', error_status_t),
565    )
566
567# 3.1.5.28 OpenPerformanceText (Opnum 32)
568class OpenPerformanceText(NDRCALL):
569    opnum = 32
570    structure = (
571       ('ServerName', PREGISTRY_SERVER_NAME),
572       ('samDesired', REGSAM),
573    )
574
575class OpenPerformanceTextResponse(NDRCALL):
576    structure = (
577       ('phKey', RPC_HKEY),
578       ('ErrorCode', error_status_t),
579    )
580
581# 3.1.5.29 OpenPerformanceNlsText (Opnum 33)
582class OpenPerformanceNlsText(NDRCALL):
583    opnum = 33
584    structure = (
585       ('ServerName', PREGISTRY_SERVER_NAME),
586       ('samDesired', REGSAM),
587    )
588
589class OpenPerformanceNlsTextResponse(NDRCALL):
590    structure = (
591       ('phKey', RPC_HKEY),
592       ('ErrorCode', error_status_t),
593    )
594
595# 3.1.5.30 BaseRegQueryMultipleValues2 (Opnum 34)
596class BaseRegQueryMultipleValues2(NDRCALL):
597    opnum = 34
598    structure = (
599       ('hKey', RPC_HKEY),
600       ('val_listIn', RVALENT_ARRAY),
601       ('num_vals', DWORD),
602       ('lpvalueBuf', PBYTE_ARRAY),
603       ('ldwTotsize', DWORD),
604    )
605
606class BaseRegQueryMultipleValues2Response(NDRCALL):
607    structure = (
608       ('val_listOut', RVALENT_ARRAY),
609       ('lpvalueBuf', PBYTE_ARRAY),
610       ('ldwRequiredSize', DWORD),
611       ('ErrorCode', error_status_t),
612    )
613
614# 3.1.5.31 BaseRegDeleteKeyEx (Opnum 35)
615class BaseRegDeleteKeyEx(NDRCALL):
616    opnum = 35
617    structure = (
618       ('hKey', RPC_HKEY),
619       ('lpSubKey', RRP_UNICODE_STRING),
620       ('AccessMask', REGSAM),
621       ('Reserved', DWORD),
622    )
623
624class BaseRegDeleteKeyExResponse(NDRCALL):
625    structure = (
626       ('ErrorCode', error_status_t),
627    )
628
629################################################################################
630# OPNUMs and their corresponding structures
631################################################################################
632OPNUMS = {
633 0 : (OpenClassesRoot, OpenClassesRootResponse),
634 1 : (OpenCurrentUser, OpenCurrentUserResponse),
635 2 : (OpenLocalMachine, OpenLocalMachineResponse),
636 3 : (OpenPerformanceData, OpenPerformanceDataResponse),
637 4 : (OpenUsers, OpenUsersResponse),
638 5 : (BaseRegCloseKey, BaseRegCloseKeyResponse),
639 6 : (BaseRegCreateKey, BaseRegCreateKeyResponse),
640 7 : (BaseRegDeleteKey, BaseRegDeleteKeyResponse),
641 8 : (BaseRegDeleteValue, BaseRegDeleteValueResponse),
642 9 : (BaseRegEnumKey, BaseRegEnumKeyResponse),
64310 : (BaseRegEnumValue, BaseRegEnumValueResponse),
64411 : (BaseRegFlushKey, BaseRegFlushKeyResponse),
64512 : (BaseRegGetKeySecurity, BaseRegGetKeySecurityResponse),
64613 : (BaseRegLoadKey, BaseRegLoadKeyResponse),
64715 : (BaseRegOpenKey, BaseRegOpenKeyResponse),
64816 : (BaseRegQueryInfoKey, BaseRegQueryInfoKeyResponse),
64917 : (BaseRegQueryValue, BaseRegQueryValueResponse),
65018 : (BaseRegReplaceKey, BaseRegReplaceKeyResponse),
65119 : (BaseRegRestoreKey, BaseRegRestoreKeyResponse),
65220 : (BaseRegSaveKey, BaseRegSaveKeyResponse),
65321 : (BaseRegSetKeySecurity, BaseRegSetKeySecurityResponse),
65422 : (BaseRegSetValue, BaseRegSetValueResponse),
65523 : (BaseRegUnLoadKey, BaseRegUnLoadKeyResponse),
65626 : (BaseRegGetVersion, BaseRegGetVersionResponse),
65727 : (OpenCurrentConfig, OpenCurrentConfigResponse),
65829 : (BaseRegQueryMultipleValues, BaseRegQueryMultipleValuesResponse),
65931 : (BaseRegSaveKeyEx, BaseRegSaveKeyExResponse),
66032 : (OpenPerformanceText, OpenPerformanceTextResponse),
66133 : (OpenPerformanceNlsText, OpenPerformanceNlsTextResponse),
66234 : (BaseRegQueryMultipleValues2, BaseRegQueryMultipleValues2Response),
66335 : (BaseRegDeleteKeyEx, BaseRegDeleteKeyExResponse),
664}
665
666################################################################################
667# HELPER FUNCTIONS
668################################################################################
669def checkNullString(string):
670    if string == NULL:
671        return string
672
673    if string[-1:] != '\x00':
674        return string + '\x00'
675    else:
676        return string
677
678def packValue(valueType, value):
679    if valueType == REG_DWORD:
680        retData = pack('<L', value)
681    elif valueType == REG_DWORD_BIG_ENDIAN:
682        retData = pack('>L', value)
683    elif valueType == REG_EXPAND_SZ:
684        try:
685            retData = value.encode('utf-16le')
686        except UnicodeDecodeError:
687            import sys
688            retData = value.decode(sys.getfilesystemencoding()).encode('utf-16le')
689    elif valueType == REG_MULTI_SZ:
690        try:
691            retData = value.encode('utf-16le')
692        except UnicodeDecodeError:
693            import sys
694            retData = value.decode(sys.getfilesystemencoding()).encode('utf-16le')
695    elif valueType == REG_QWORD:
696        retData = pack('<Q', value)
697    elif valueType == REG_QWORD_LITTLE_ENDIAN:
698        retData = pack('>Q', value)
699    elif valueType == REG_SZ:
700        try:
701            retData = value.encode('utf-16le')
702        except UnicodeDecodeError:
703            import sys
704            retData = value.decode(sys.getfilesystemencoding()).encode('utf-16le')
705    else:
706        retData = value
707
708    return retData
709
710def unpackValue(valueType, value):
711    if valueType == REG_DWORD:
712        retData = unpack('<L', ''.join(value))[0]
713    elif valueType == REG_DWORD_BIG_ENDIAN:
714        retData = unpack('>L', ''.join(value))[0]
715    elif valueType == REG_EXPAND_SZ:
716        retData = ''.join(value).decode('utf-16le')
717    elif valueType == REG_MULTI_SZ:
718        retData = ''.join(value).decode('utf-16le')
719    elif valueType == REG_QWORD:
720        retData = unpack('<Q', ''.join(value))[0]
721    elif valueType == REG_QWORD_LITTLE_ENDIAN:
722        retData = unpack('>Q', ''.join(value))[0]
723    elif valueType == REG_SZ:
724        retData = ''.join(value).decode('utf-16le')
725    else:
726        retData = ''.join(value)
727
728    return retData
729
730def hOpenClassesRoot(dce, samDesired = MAXIMUM_ALLOWED):
731    request = OpenClassesRoot()
732    request['ServerName'] = NULL
733    request['samDesired'] = samDesired
734    return dce.request(request)
735
736def hOpenCurrentUser(dce, samDesired = MAXIMUM_ALLOWED):
737    request = OpenCurrentUser()
738    request['ServerName'] = NULL
739    request['samDesired'] = samDesired
740    return dce.request(request)
741
742def hOpenLocalMachine(dce, samDesired = MAXIMUM_ALLOWED):
743    request = OpenLocalMachine()
744    request['ServerName'] = NULL
745    request['samDesired'] = samDesired
746    return dce.request(request)
747
748def hOpenPerformanceData(dce, samDesired = MAXIMUM_ALLOWED):
749    request = OpenPerformanceData()
750    request['ServerName'] = NULL
751    request['samDesired'] = samDesired
752    return dce.request(request)
753
754def hOpenUsers(dce, samDesired = MAXIMUM_ALLOWED):
755    request = OpenUsers()
756    request['ServerName'] = NULL
757    request['samDesired'] = samDesired
758    return dce.request(request)
759
760def hBaseRegCloseKey(dce, hKey):
761    request = BaseRegCloseKey()
762    request['hKey'] = hKey
763    return dce.request(request)
764
765def hBaseRegCreateKey(dce, hKey, lpSubKey, lpClass = NULL, dwOptions = 0x00000001, samDesired = MAXIMUM_ALLOWED, lpSecurityAttributes = NULL, lpdwDisposition = REG_CREATED_NEW_KEY):
766    request = BaseRegCreateKey()
767    request['hKey'] = hKey
768    request['lpSubKey'] = checkNullString(lpSubKey)
769    request['lpClass'] = checkNullString(lpClass)
770    request['dwOptions'] = dwOptions
771    request['samDesired'] = samDesired
772    if lpSecurityAttributes == NULL:
773        request['lpSecurityAttributes']['RpcSecurityDescriptor']['lpSecurityDescriptor'] = NULL
774    else:
775        request['lpSecurityAttributes'] = lpSecurityAttributes
776    request['lpdwDisposition'] = lpdwDisposition
777
778    return dce.request(request)
779
780def hBaseRegDeleteKey(dce, hKey, lpSubKey):
781    request = BaseRegDeleteKey()
782    request['hKey'] = hKey
783    request['lpSubKey'] = checkNullString(lpSubKey)
784    return dce.request(request)
785
786def hBaseRegEnumKey(dce, hKey, dwIndex, lpftLastWriteTime = NULL):
787    request = BaseRegEnumKey()
788    request['hKey'] = hKey
789    request['dwIndex'] = dwIndex
790    request.fields['lpNameIn'].fields['MaximumLength'] = 1024
791    request.fields['lpNameIn'].fields['Data'].fields['Data'].fields['MaximumCount'] = 1024/2
792    request['lpClassIn'] = ' '* 64
793    request['lpftLastWriteTime'] = lpftLastWriteTime
794
795    return dce.request(request)
796
797def hBaseRegEnumValue(dce, hKey, dwIndex, dataLen=256):
798    request = BaseRegEnumValue()
799    request['hKey'] = hKey
800    request['dwIndex'] = dwIndex
801    retries = 1
802
803    # We need to be aware the size might not be enough, so let's catch ERROR_MORE_DATA exception
804    while True:
805        try:
806            # Only the maximum length field of the lpValueNameIn is used to determine the buffer length to be allocated
807            # by the service. Specify a string with a zero length but maximum length set to the largest buffer size
808            # needed to hold the value names.
809            request.fields['lpValueNameIn'].fields['MaximumLength'] = dataLen*2
810            request.fields['lpValueNameIn'].fields['Data'].fields['Data'].fields['MaximumCount'] = dataLen
811
812            request['lpData'] = ' ' * dataLen
813            request['lpcbData'] = dataLen
814            request['lpcbLen'] = dataLen
815            resp = dce.request(request)
816        except DCERPCSessionError, e:
817            if retries > 1:
818                LOG.debug('Too many retries when calling hBaseRegEnumValue, aborting')
819                raise
820            if e.get_error_code() == system_errors.ERROR_MORE_DATA:
821                # We need to adjust the size
822                retries +=1
823                dataLen = e.get_packet()['lpcbData']
824                continue
825            else:
826                raise
827        else:
828            break
829
830    return resp
831
832def hBaseRegFlushKey(dce, hKey):
833    request = BaseRegFlushKey()
834    request['hKey'] = hKey
835    return dce.request(request)
836
837def hBaseRegGetKeySecurity(dce, hKey, securityInformation = OWNER_SECURITY_INFORMATION ):
838    request = BaseRegGetKeySecurity()
839    request['hKey'] = hKey
840    request['SecurityInformation'] = securityInformation
841    request['pRpcSecurityDescriptorIn']['lpSecurityDescriptor'] = NULL
842    request['pRpcSecurityDescriptorIn']['cbInSecurityDescriptor'] = 1024
843
844    return dce.request(request)
845
846def hBaseRegLoadKey(dce, hKey, lpSubKey, lpFile):
847    request = BaseRegLoadKey()
848    request['hKey'] = hKey
849    request['lpSubKey'] = checkNullString(lpSubKey)
850    request['lpFile'] = checkNullString(lpFile)
851    return dce.request(request)
852
853def hBaseRegUnLoadKey(dce, hKey, lpSubKey):
854    request = BaseRegUnLoadKey()
855    request['hKey'] = hKey
856    request['lpSubKey'] = checkNullString(lpSubKey)
857    return dce.request(request)
858
859def hBaseRegOpenKey(dce, hKey, lpSubKey, dwOptions=0x00000001, samDesired = MAXIMUM_ALLOWED):
860    request = BaseRegOpenKey()
861    request['hKey'] = hKey
862    request['lpSubKey'] = checkNullString(lpSubKey)
863    request['dwOptions'] = dwOptions
864    request['samDesired'] = samDesired
865    return dce.request(request)
866
867def hBaseRegQueryInfoKey(dce, hKey):
868    request = BaseRegQueryInfoKey()
869    request['hKey'] = hKey
870    # Not the cleanest way, but oh well
871    # Plus, Windows XP needs MaximumCount also set
872    request.fields['lpClassIn'].fields['MaximumLength'] = 1024
873    request.fields['lpClassIn'].fields['Data'].fields['Data'].fields['MaximumCount'] = 1024/2
874    return dce.request(request)
875
876def hBaseRegQueryValue(dce, hKey, lpValueName, dataLen=512):
877    request = BaseRegQueryValue()
878    request['hKey'] = hKey
879    request['lpValueName'] = checkNullString(lpValueName)
880    retries = 1
881
882    # We need to be aware the size might not be enough, so let's catch ERROR_MORE_DATA exception
883    while True:
884        try:
885            request['lpData'] = ' ' * dataLen
886            request['lpcbData'] = dataLen
887            request['lpcbLen'] = dataLen
888            resp = dce.request(request)
889        except DCERPCSessionError, e:
890            if retries > 1:
891                LOG.debug('Too many retries when calling hBaseRegQueryValue, aborting')
892                raise
893            if e.get_error_code() == system_errors.ERROR_MORE_DATA:
894                # We need to adjust the size
895                dataLen = e.get_packet()['lpcbData']
896                continue
897            else:
898                raise
899        else:
900            break
901
902    # Returns
903    # ( dataType, data )
904    return resp['lpType'], unpackValue(resp['lpType'], resp['lpData'])
905
906def hBaseRegReplaceKey(dce, hKey, lpSubKey, lpNewFile, lpOldFile):
907    request = BaseRegReplaceKey()
908    request['hKey'] = hKey
909    request['lpSubKey'] = checkNullString(lpSubKey)
910    request['lpNewFile'] = checkNullString(lpNewFile)
911    request['lpOldFile'] = checkNullString(lpOldFile)
912    return dce.request(request)
913
914def hBaseRegRestoreKey(dce, hKey, lpFile, flags=REG_REFRESH_HIVE):
915    request = BaseRegRestoreKey()
916    request['hKey'] = hKey
917    request['lpFile'] = checkNullString(lpFile)
918    request['Flags'] = flags
919    return dce.request(request)
920
921def hBaseRegSaveKey(dce, hKey, lpFile, pSecurityAttributes = NULL):
922    request = BaseRegSaveKey()
923    request['hKey'] = hKey
924    request['lpFile'] = checkNullString(lpFile)
925    request['pSecurityAttributes'] = pSecurityAttributes
926    return dce.request(request)
927
928def hBaseRegSetValue(dce, hKey, lpValueName, dwType, lpData):
929    request = BaseRegSetValue()
930    request['hKey'] = hKey
931    request['lpValueName'] = checkNullString(lpValueName)
932    request['dwType'] = dwType
933    request['lpData'] = packValue(dwType,lpData)
934    request['cbData'] = len(request['lpData'])
935    return dce.request(request)
936
937def hBaseRegGetVersion(dce, hKey):
938    request = BaseRegGetVersion()
939    request['hKey'] = hKey
940    return dce.request(request)
941
942def hOpenCurrentConfig(dce, samDesired = MAXIMUM_ALLOWED):
943    request = OpenCurrentConfig()
944    request['ServerName'] = NULL
945    request['samDesired'] = samDesired
946    return dce.request(request)
947
948def hBaseRegQueryMultipleValues(dce, hKey, val_listIn):
949    # ToDo, check the result to see whether we need to
950    # have a bigger buffer for the data to receive
951    request = BaseRegQueryMultipleValues()
952    request['hKey'] = hKey
953
954    for item in  val_listIn:
955        itemn = RVALENT()
956        itemn['ve_valuename'] = checkNullString(item['ValueName'])
957        itemn['ve_valuelen'] = len(itemn['ve_valuename'])
958        itemn['ve_valueptr'] = NULL
959        itemn['ve_type'] = item['ValueType']
960        request['val_listIn'].append(itemn)
961
962    request['num_vals'] = len(request['val_listIn'])
963    request['lpvalueBuf'] = list(' '*128)
964    request['ldwTotsize'] = 128
965
966    resp = dce.request(request)
967    retVal = list()
968    for item in resp['val_listOut']:
969        itemn = dict()
970        itemn['ValueName'] = item['ve_valuename']
971        itemn['ValueData'] = unpackValue(item['ve_type'], resp['lpvalueBuf'][item['ve_valueptr'] : item['ve_valueptr']+item['ve_valuelen']])
972        retVal.append(itemn)
973
974    return retVal
975
976def hBaseRegSaveKeyEx(dce, hKey, lpFile, pSecurityAttributes = NULL, flags=1):
977    request = BaseRegSaveKeyEx()
978    request['hKey'] = hKey
979    request['lpFile'] = checkNullString(lpFile)
980    request['pSecurityAttributes'] = pSecurityAttributes
981    request['Flags'] = flags
982    return dce.request(request)
983
984def hOpenPerformanceText(dce, samDesired = MAXIMUM_ALLOWED):
985    request = OpenPerformanceText()
986    request['ServerName'] = NULL
987    request['samDesired'] = samDesired
988    return dce.request(request)
989
990def hOpenPerformanceNlsText(dce, samDesired = MAXIMUM_ALLOWED):
991    request = OpenPerformanceNlsText()
992    request['ServerName'] = NULL
993    request['samDesired'] = samDesired
994    return dce.request(request)
995
996def hBaseRegDeleteValue(dce, hKey, lpValueName):
997    request = BaseRegDeleteValue()
998    request['hKey'] = hKey
999    request['lpValueName'] = checkNullString(lpValueName)
1000    return dce.request(request)
1001
1002