cryptopals/p34.py at master · AdrS/cryptopals · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
from p33 import *
from util import *
from os import urandom
import socket
import aes
import struct
import sys
from netutil import *

def dh_connect(addr, p=nist_p, g=nist_g):
	'''sets up tcp connection to host and does dh key exchange'''
	a, A = gen_key_pair(p, g)
	con = socket.socket()
	con.connect(addr)
	print 'Connecting to %s:%d' % addr
	send_bigint(con, p)
	send_bigint(con, g)
	send_bigint(con, A)
	B = recv_bigint(con)
	s = get_common_key(B, a, p)
	key = sha1sum(bigint_to_bytes(s))[:16]
	print 'Shared key: ', key.encode('hex')
	return con, key

def send_ct(con, msg, key):
	'''sends <len> <aes cbc encrypted msg>'''
	iv = urandom(16)
	return send_len_payload(con, aes.encrypt_cbc(msg, key, iv))

def recv_ct(con, key):
	'''reads <len> <aes cbc encrypted msg>'''
	return aes.decrypt_cbc(recv_len_payload(con), key)

def echo_server_handler(con, addr):
	p = recv_bigint(con)
	g = recv_bigint(con)
	A = recv_bigint(con)
	if g >= p or A > p:
		print 'received bad parameters'
		con.send('error: bad parameters')
		con.close()
		return
	b, B = gen_key_pair(p, g)
	send_bigint(con, B)
	s = get_common_key(A, b, p)
	key = sha1sum(bigint_to_bytes(s))[:16]
	print 'Established shared key:', key.encode('hex')

	msg = recv_ct(con, key)

	print 'Received: ', msg

	#echo message back
	send_ct(con, msg, key)
	con.close()

def ping_server(addr):
	con, key = dh_connect(addr)
	send_ct(con, 'Ping from Alice', key)
	print 'Received', recv_ct(con, key)
	con.close()

def mitm_server_handler(conA, addrA, addrB):
	#start DH with Alice
	#A -> M p, g, A
	p = recv_bigint(conA)
	g = recv_bigint(conA)
	A = recv_bigint(conA)
	if g >= p or A >= p:
		print 'received bad parameters'
		conA.send('error: bad parameters')
		conA.close()
		return

	#do DH with Bob
	#M -> B p, g, p
	conB = socket.socket()
	conB.connect(addrB)
	send_bigint(conB, p)
	send_bigint(conB, g)
	send_bigint(conB, p)

	#B -> M B
	B = recv_bigint(conB)

	#sB = p^b = 0 (mod p)
	key = sha1sum(bigint_to_bytes(0))[:16]

	print 'Finished DH with Bob'

	#finish DH with Alice
	#M -> A p
	send_bigint(conA, p)

	#sA = p^a = 0 (mod p)

	print 'Finished DH with Alice'

	#NOTE: Alice and Bob always end up with the same predictable
	#	session key because of the MITM

	#read ping from Alice
	msgA = recv_ct(conA, key)

	print 'Received ping from Alice: ', msgA

	#forward ping to Bob
	send_ct(conB, msgA, key)

	#read pong from Bob
	msgB = recv_ct(conB, key)
	conB.close()

	print 'Received pong from Bob: ', msgB

	#forward pong to Alice
	send_ct(conA, msgB, key)
	conA.close()

def usage(argv):
	print ''''usage:
	%s <port>                   run echo server on port
	%s <host> <port>            ping echo server at address
	%s <portl> <host> <ports>   act as mitm listening on <portl> and
	                            forwarding traffic to server at <host> <ports>
	''' % (argv[0], argv[0], argv[0])
	sys.exit(1)

if __name__ == '__main__':
	if len(sys.argv) == 2:
		serve_forever(int(sys.argv[1]), echo_server_handler)
	elif len(sys.argv) == 3:
		ping_server((sys.argv[1], int(sys.argv[2])))
	elif len(sys.argv) == 4:
		addrB = (sys.argv[2], int(sys.argv[3]))
		handler = lambda conA, addrA: mitm_server_handler(conA, addrA, addrB)
		serve_forever(int(sys.argv[1]), handler)
	else:
		usage(sys.argv)