An stunnel like program that utilizes the Noise protocol.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

896 lines
24 KiB

  1. from cryptography.hazmat.backends import default_backend
  2. from cryptography.hazmat.primitives import hashes
  3. from cryptography.hazmat.primitives import serialization
  4. from cryptography.hazmat.primitives.asymmetric import x448
  5. from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
  6. from cryptography.hazmat.primitives.kdf.hkdf import HKDF
  7. from cryptography.hazmat.primitives.serialization import load_pem_private_key
  8. from noise.connection import NoiseConnection, Keypair
  9. #import tracemalloc; tracemalloc.start()
  10. import argparse
  11. import asyncio
  12. import base64
  13. import os.path
  14. import shutil
  15. import socket
  16. import sys
  17. import tempfile
  18. import time
  19. import threading
  20. import unittest
  21. _backend = default_backend()
  22. def loadprivkey(fname):
  23. with open(fname, encoding='ascii') as fp:
  24. data = fp.read().encode('ascii')
  25. key = load_pem_private_key(data, password=None, backend=default_backend())
  26. return key
  27. def loadprivkeyraw(fname):
  28. key = loadprivkey(fname)
  29. enc = serialization.Encoding.Raw
  30. privformat = serialization.PrivateFormat.Raw
  31. encalgo = serialization.NoEncryption()
  32. return key.private_bytes(encoding=enc, format=privformat, encryption_algorithm=encalgo)
  33. def loadpubkeyraw(fname):
  34. with open(fname, encoding='ascii') as fp:
  35. lines = fp.readlines()
  36. # XXX
  37. #self.assertEqual(len(lines), 1)
  38. keytype, keyvalue = lines[0].split()
  39. if keytype != 'ntun-x448':
  40. raise RuntimeError
  41. return base64.urlsafe_b64decode(keyvalue)
  42. def genkeypair():
  43. '''Generates a keypair, and returns a tuple of (public, private).
  44. They are encoded as raw bytes, and sutible for use w/ Noise.'''
  45. key = x448.X448PrivateKey.generate()
  46. enc = serialization.Encoding.Raw
  47. pubformat = serialization.PublicFormat.Raw
  48. privformat = serialization.PrivateFormat.Raw
  49. encalgo = serialization.NoEncryption()
  50. pub = key.public_key().public_bytes(encoding=enc, format=pubformat)
  51. priv = key.private_bytes(encoding=enc, format=privformat, encryption_algorithm=encalgo)
  52. return pub, priv
  53. def _makefut(obj):
  54. loop = asyncio.get_running_loop()
  55. fut = loop.create_future()
  56. fut.set_result(obj)
  57. return fut
  58. def _makeunix(path):
  59. '''Make a properly formed unix path socket string.'''
  60. return 'unix:%s' % path
  61. def _parsesockstr(sockstr):
  62. proto, rem = sockstr.split(':', 1)
  63. return proto, rem
  64. async def connectsockstr(sockstr):
  65. proto, rem = _parsesockstr(sockstr)
  66. reader, writer = await asyncio.open_unix_connection(rem)
  67. return reader, writer
  68. async def listensockstr(sockstr, cb):
  69. '''Wrapper for asyncio.start_x_server.
  70. The format of sockstr is: 'proto:param=value[,param2=value2]'.
  71. If the proto has a default parameter, the value can be used
  72. directly, like: 'proto:value'. This is only allowed when the
  73. value can unambiguously be determined not to be a param.
  74. The characters that define 'param' must be all lower case ascii
  75. characters and may contain an underscore. The first character
  76. must not be and underscore.
  77. Supported protocols:
  78. unix:
  79. Default parameter is path.
  80. The path parameter specifies the path to the
  81. unix domain socket. The path MUST start w/ a
  82. slash if it is used as a default parameter.
  83. '''
  84. proto, rem = _parsesockstr(sockstr)
  85. server = await asyncio.start_unix_server(cb, path=rem)
  86. return server
  87. # !!python makemessagelengths.py
  88. _handshakelens = \
  89. [72, 72, 88]
  90. def _genciphfun(hash, ad):
  91. hkdf = HKDF(algorithm=hashes.SHA256(), length=32,
  92. salt=b'asdoifjsldkjdsf', info=ad, backend=_backend)
  93. key = hkdf.derive(hash)
  94. cipher = Cipher(algorithms.AES(key), modes.ECB(),
  95. backend=_backend)
  96. enctor = cipher.encryptor()
  97. def encfun(data):
  98. # Returns the two bytes for length
  99. val = len(data)
  100. encbytes = enctor.update(data[:16])
  101. mask = int.from_bytes(encbytes[:2], byteorder='big') & 0xff
  102. return (val ^ mask).to_bytes(length=2, byteorder='big')
  103. def decfun(data):
  104. # takes off the data and returns the total
  105. # length
  106. val = int.from_bytes(data[:2], byteorder='big')
  107. encbytes = enctor.update(data[2:2 + 16])
  108. mask = int.from_bytes(encbytes[:2], byteorder='big') & 0xff
  109. return val ^ mask
  110. return encfun, decfun
  111. async def NoiseForwarder(mode, encrdrwrr, ptpairfun, priv_key, pub_key=None):
  112. '''A function that forwards data between the plain text pair of
  113. streams to the encrypted session.
  114. The mode paramater must be one of 'init' or 'resp' for initiator
  115. and responder.
  116. The encrdrwrr is an await object that will return a tunle of the
  117. reader and writer streams for the encrypted side of the
  118. connection.
  119. The ptpairfun parameter is a function that will be passed the
  120. public key bytes for the remote client. This can be used to
  121. both validate that the correct client is connecting, and to
  122. pass back the correct plain text reader/writer objects that
  123. match the provided static key. The function must be an async
  124. function.
  125. In the case of the initiator, pub_key must be provided and will
  126. be used to authenticate the responder side of the connection.
  127. The priv_key parameter is used to authenticate this side of the
  128. session.
  129. Both priv_key and pub_key parameters must be 56 bytes. For example,
  130. the pair that is returned by genkeypair.
  131. '''
  132. rdr, wrr = await encrdrwrr
  133. proto = NoiseConnection.from_name(b'Noise_XK_448_ChaChaPoly_SHA256')
  134. proto.set_keypair_from_private_bytes(Keypair.STATIC, priv_key)
  135. if pub_key is not None:
  136. proto.set_keypair_from_public_bytes(Keypair.REMOTE_STATIC,
  137. pub_key)
  138. if mode == 'resp':
  139. proto.set_as_responder()
  140. proto.start_handshake()
  141. proto.read_message(await rdr.readexactly(_handshakelens[0]))
  142. wrr.write(proto.write_message())
  143. proto.read_message(await rdr.readexactly(_handshakelens[2]))
  144. elif mode == 'init':
  145. proto.set_as_initiator()
  146. proto.start_handshake()
  147. wrr.write(proto.write_message())
  148. proto.read_message(await rdr.readexactly(_handshakelens[1]))
  149. wrr.write(proto.write_message())
  150. if not proto.handshake_finished: # pragma: no cover
  151. raise RuntimeError('failed to finish handshake')
  152. reader, writer = await ptpairfun(getattr(proto.get_keypair(
  153. Keypair.REMOTE_STATIC), 'public_bytes', None))
  154. # generate the keys for lengths
  155. # XXX - get_handshake_hash is probably not the best option, but
  156. # this is only to obscure lengths, it is not required to be secure
  157. # as the underlying NoiseProtocol securely validates everything.
  158. # It is marginally useful as writing patterns likely expose the
  159. # true length. Adding padding could marginally help w/ this.
  160. if mode == 'resp':
  161. _, declenfun = _genciphfun(proto.get_handshake_hash(), b'toresp')
  162. enclenfun, _ = _genciphfun(proto.get_handshake_hash(), b'toinit')
  163. elif mode == 'init':
  164. enclenfun, _ = _genciphfun(proto.get_handshake_hash(), b'toresp')
  165. _, declenfun = _genciphfun(proto.get_handshake_hash(), b'toinit')
  166. async def decses():
  167. try:
  168. while True:
  169. try:
  170. msg = await rdr.readexactly(2 + 16)
  171. except asyncio.streams.IncompleteReadError:
  172. if rdr.at_eof():
  173. return 'dec'
  174. tlen = declenfun(msg)
  175. rmsg = await rdr.readexactly(tlen - 16)
  176. tmsg = msg[2:] + rmsg
  177. writer.write(proto.decrypt(tmsg))
  178. await writer.drain()
  179. #except:
  180. # import traceback
  181. # traceback.print_exc()
  182. # raise
  183. finally:
  184. try:
  185. writer.write_eof()
  186. except OSError as e:
  187. if e.errno != 57:
  188. raise
  189. async def encses():
  190. try:
  191. while True:
  192. # largest message
  193. ptmsg = await reader.read(65535 - 16)
  194. if not ptmsg:
  195. # eof
  196. return 'enc'
  197. encmsg = proto.encrypt(ptmsg)
  198. wrr.write(enclenfun(encmsg))
  199. wrr.write(encmsg)
  200. await wrr.drain()
  201. #except:
  202. # import traceback
  203. # traceback.print_exc()
  204. # raise
  205. finally:
  206. wrr.write_eof()
  207. return await asyncio.gather(decses(), encses())
  208. # https://stackoverflow.com/questions/23033939/how-to-test-python-3-4-asyncio-code
  209. # Slightly modified to timeout and to print trace back when canceled.
  210. # This makes it easier to figure out what "froze".
  211. def async_test(f):
  212. def wrapper(*args, **kwargs):
  213. async def tbcapture():
  214. try:
  215. return await f(*args, **kwargs)
  216. except asyncio.CancelledError as e:
  217. # if we are going to be cancelled, print out a tb
  218. import traceback
  219. traceback.print_exc()
  220. raise
  221. loop = asyncio.get_event_loop()
  222. # timeout after 4 seconds
  223. loop.run_until_complete(asyncio.wait_for(tbcapture(), 4))
  224. return wrapper
  225. class Tests_misc(unittest.TestCase):
  226. def test_listensockstr(self):
  227. # XXX write test
  228. pass
  229. def test_genciphfun(self):
  230. enc, dec = _genciphfun(b'0' * 32, b'foobar')
  231. msg = b'this is a bunch of data'
  232. tb = enc(msg)
  233. self.assertEqual(len(msg), dec(tb + msg))
  234. for i in [ 20, 1384, 64000, 23839, 65535 ]:
  235. msg = os.urandom(i)
  236. self.assertEqual(len(msg), dec(enc(msg) + msg))
  237. def cmd_client(args):
  238. privkey = loadprivkeyraw(args.clientkey)
  239. pubkey = loadpubkeyraw(args.servkey)
  240. async def runnf(rdr, wrr):
  241. encpair = asyncio.create_task(connectsockstr(args.clienttarget))
  242. a = await NoiseForwarder('init',
  243. encpair, lambda x: _makefut((rdr, wrr)),
  244. priv_key=privkey, pub_key=pubkey)
  245. # Setup client listener
  246. ssock = listensockstr(args.clientlisten, runnf)
  247. loop = asyncio.get_event_loop()
  248. obj = loop.run_until_complete(ssock)
  249. loop.run_until_complete(obj.serve_forever())
  250. def cmd_server(args):
  251. privkey = loadprivkeyraw(args.servkey)
  252. async def runnf(rdr, wrr):
  253. ptpairfun = asyncio.create_task(connectsockstr(args.servtarget))
  254. a = await NoiseForwarder('resp',
  255. _makefut((rdr, wrr)), lambda x: ptpairfun,
  256. priv_key=privkey)
  257. # Setup server listener
  258. ssock = listensockstr(args.servlisten, runnf)
  259. loop = asyncio.get_event_loop()
  260. obj = loop.run_until_complete(ssock)
  261. loop.run_until_complete(obj.serve_forever())
  262. def cmd_genkey(args):
  263. keypair = genkeypair()
  264. key = x448.X448PrivateKey.generate()
  265. # public key part
  266. enc = serialization.Encoding.Raw
  267. pubformat = serialization.PublicFormat.Raw
  268. pub = key.public_key().public_bytes(encoding=enc, format=pubformat)
  269. try:
  270. fname = args.fname + '.pub'
  271. with open(fname, 'x', encoding='ascii') as fp:
  272. print('ntun-x448', base64.urlsafe_b64encode(pub).decode('ascii'), file=fp)
  273. except FileExistsError:
  274. print('failed to create %s, file exists.' % fname, file=sys.stderr)
  275. sys.exit(2)
  276. enc = serialization.Encoding.PEM
  277. format = serialization.PrivateFormat.PKCS8
  278. encalgo = serialization.NoEncryption()
  279. with open(args.fname, 'x', encoding='ascii') as fp:
  280. fp.write(key.private_bytes(encoding=enc, format=format, encryption_algorithm=encalgo).decode('ascii'))
  281. def main():
  282. parser = argparse.ArgumentParser()
  283. subparsers = parser.add_subparsers(title='subcommands', description='valid subcommands', help='additional help')
  284. parser_gk = subparsers.add_parser('genkey', help='generate keys')
  285. parser_gk.add_argument('fname', type=str, help='file name for the key')
  286. parser_gk.set_defaults(func=cmd_genkey)
  287. parser_serv = subparsers.add_parser('server', help='run a server')
  288. parser_serv.add_argument('-c', action='append', type=str, help='file of authorized client keys, or a .pub file')
  289. parser_serv.add_argument('servkey', type=str, help='file name for the server key')
  290. parser_serv.add_argument('servlisten', type=str, help='Connection that the server listens on')
  291. parser_serv.add_argument('servtarget', type=str, help='Connection that the server connects to')
  292. parser_serv.set_defaults(func=cmd_server)
  293. parser_client = subparsers.add_parser('client', help='run a client')
  294. parser_client.add_argument('clientkey', type=str, help='file name for the client private key')
  295. parser_client.add_argument('servkey', type=str, help='file name for the server public key')
  296. parser_client.add_argument('clientlisten', type=str, help='Connection that the client listens on')
  297. parser_client.add_argument('clienttarget', type=str, help='Connection that the client connects to')
  298. parser_client.set_defaults(func=cmd_client)
  299. args = parser.parse_args()
  300. try:
  301. fun = args.func
  302. except AttributeError:
  303. parser.print_usage()
  304. sys.exit(5)
  305. fun(args)
  306. if __name__ == '__main__': # pragma: no cover
  307. main()
  308. def _asyncsockpair():
  309. '''Create a pair of sockets that are bound to each other.
  310. The function will return a tuple of two coroutine's, that
  311. each, when await'ed upon, will return the reader/writer pair.'''
  312. socka, sockb = socket.socketpair()
  313. return asyncio.open_connection(sock=socka), \
  314. asyncio.open_connection(sock=sockb)
  315. async def _awaitfile(fname):
  316. while not os.path.exists(fname):
  317. await asyncio.sleep(.01)
  318. return True
  319. class TestMain(unittest.TestCase):
  320. def setUp(self):
  321. # setup temporary directory
  322. d = os.path.realpath(tempfile.mkdtemp())
  323. self.basetempdir = d
  324. self.tempdir = os.path.join(d, 'subdir')
  325. os.mkdir(self.tempdir)
  326. # Generate key pairs
  327. self.server_key_pair = genkeypair()
  328. self.client_key_pair = genkeypair()
  329. os.chdir(self.tempdir)
  330. def tearDown(self):
  331. #print('td:', time.time())
  332. shutil.rmtree(self.basetempdir)
  333. self.tempdir = None
  334. @async_test
  335. async def test_noargs(self):
  336. proc = await self.run_with_args()
  337. await proc.wait()
  338. # XXX - not checking error message
  339. # And that it exited w/ the correct code
  340. self.assertEqual(proc.returncode, 5)
  341. def run_with_args(self, *args, pipes=True):
  342. kwargs = {}
  343. if pipes:
  344. kwargs.update(dict(
  345. stdout=asyncio.subprocess.PIPE,
  346. stderr=asyncio.subprocess.PIPE))
  347. return asyncio.create_subprocess_exec(sys.executable,
  348. # XXX - figure out how to add coverage data on these runs
  349. #'-m', 'coverage', 'run', '-p',
  350. __file__, *args, **kwargs)
  351. async def genkey(self, name):
  352. proc = await self.run_with_args('genkey', name, pipes=False)
  353. await proc.wait()
  354. self.assertEqual(proc.returncode, 0)
  355. @async_test
  356. async def test_loadpubkey(self):
  357. keypath = os.path.join(self.tempdir, 'loadpubkeytest')
  358. await self.genkey(keypath)
  359. privkey = loadprivkey(keypath)
  360. enc = serialization.Encoding.Raw
  361. pubformat = serialization.PublicFormat.Raw
  362. pubkeybytes = privkey.public_key().public_bytes(encoding=enc, format=pubformat)
  363. pubkey = loadpubkeyraw(keypath + '.pub')
  364. self.assertEqual(pubkeybytes, pubkey)
  365. privrawkey = loadprivkeyraw(keypath)
  366. enc = serialization.Encoding.Raw
  367. privformat = serialization.PrivateFormat.Raw
  368. encalgo = serialization.NoEncryption()
  369. rprivrawkey = privkey.private_bytes(encoding=enc, format=privformat, encryption_algorithm=encalgo)
  370. self.assertEqual(rprivrawkey, privrawkey)
  371. @async_test
  372. async def test_end2end(self):
  373. # Generate necessar keys
  374. servkeypath = os.path.join(self.tempdir, 'server_key')
  375. await self.genkey(servkeypath)
  376. clientkeypath = os.path.join(self.tempdir, 'client_key')
  377. await self.genkey(clientkeypath)
  378. await asyncio.sleep(.1)
  379. #import pdb; pdb.set_trace()
  380. # forwards connectsion to this socket (created by client)
  381. ptclientpath = os.path.join(self.tempdir, 'incclient.sock')
  382. ptclientstr = _makeunix(ptclientpath)
  383. # this is the socket server listen to
  384. incservpath = os.path.join(self.tempdir, 'incserv.sock')
  385. incservstr = _makeunix(incservpath)
  386. # to this socket, opened by server
  387. servtargpath = os.path.join(self.tempdir, 'servtarget.sock')
  388. servtargstr = _makeunix(servtargpath)
  389. # Setup server target listener
  390. ptsock = []
  391. ptsockevent = asyncio.Event()
  392. def ptsockaccept(reader, writer, ptsock=ptsock):
  393. ptsock.append((reader, writer))
  394. ptsockevent.set()
  395. # Bind to pt listener
  396. lsock = await listensockstr(servtargstr, ptsockaccept)
  397. # Startup the server
  398. server = await self.run_with_args('server',
  399. '-c', clientkeypath + '.pub',
  400. servkeypath, incservstr, servtargstr,
  401. pipes=False)
  402. # Startup the client
  403. client = await self.run_with_args('client',
  404. clientkeypath, servkeypath + '.pub', ptclientstr, incservstr,
  405. pipes=False)
  406. # wait for server target to be created
  407. await _awaitfile(servtargpath)
  408. # wait for server to start
  409. await _awaitfile(incservpath)
  410. # wait for client to start
  411. await _awaitfile(ptclientpath)
  412. await asyncio.sleep(.1)
  413. # Connect to the client
  414. reader, writer = await connectsockstr(ptclientstr)
  415. # send a message
  416. ptmsg = b'this is a message for testing'
  417. writer.write(ptmsg)
  418. # make sure that we got the conenction
  419. await ptsockevent.wait()
  420. # get the connection
  421. endrdr, endwrr = ptsock[0]
  422. # make sure we can read back what we sent
  423. self.assertEqual(ptmsg, await endrdr.readexactly(len(ptmsg)))
  424. # test some additional messages
  425. for i in [ 129, 1287, 28792, 129872 ]:
  426. # in on direction
  427. msg = os.urandom(i)
  428. writer.write(msg)
  429. self.assertEqual(msg, await endrdr.readexactly(len(msg)))
  430. # and the other
  431. endwrr.write(msg)
  432. self.assertEqual(msg, await reader.readexactly(len(msg)))
  433. @async_test
  434. async def test_genkey(self):
  435. # that it can generate a key
  436. proc = await self.run_with_args('genkey', 'somefile')
  437. await proc.wait()
  438. #print(await proc.communicate())
  439. self.assertEqual(proc.returncode, 0)
  440. with open('somefile.pub', encoding='ascii') as fp:
  441. lines = fp.readlines()
  442. self.assertEqual(len(lines), 1)
  443. keytype, keyvalue = lines[0].split()
  444. self.assertEqual(keytype, 'ntun-x448')
  445. key = x448.X448PublicKey.from_public_bytes(base64.urlsafe_b64decode(keyvalue))
  446. key = loadprivkey('somefile')
  447. self.assertIsInstance(key, x448.X448PrivateKey)
  448. # that a second call fails
  449. proc = await self.run_with_args('genkey', 'somefile')
  450. await proc.wait()
  451. stdoutdata, stderrdata = await proc.communicate()
  452. self.assertFalse(stdoutdata)
  453. self.assertEqual(b'failed to create somefile.pub, file exists.\n', stderrdata)
  454. # And that it exited w/ the correct code
  455. self.assertEqual(proc.returncode, 2)
  456. class TestNoiseFowarder(unittest.TestCase):
  457. def setUp(self):
  458. # setup temporary directory
  459. d = os.path.realpath(tempfile.mkdtemp())
  460. self.basetempdir = d
  461. self.tempdir = os.path.join(d, 'subdir')
  462. os.mkdir(self.tempdir)
  463. # Generate key pairs
  464. self.server_key_pair = genkeypair()
  465. self.client_key_pair = genkeypair()
  466. def tearDown(self):
  467. shutil.rmtree(self.basetempdir)
  468. self.tempdir = None
  469. @async_test
  470. async def test_clientkeymissmatch(self):
  471. # generate a key that is incorrect
  472. wrongclient_key_pair = genkeypair()
  473. # the secure socket
  474. clssockapair, clssockbpair = _asyncsockpair()
  475. reader, writer = await clssockapair
  476. async def wrongkey(v):
  477. raise ValueError('no key matches')
  478. # create the server
  479. servnf = asyncio.create_task(NoiseForwarder('resp',
  480. clssockbpair, wrongkey,
  481. priv_key=self.server_key_pair[1]))
  482. # Create client
  483. proto = NoiseConnection.from_name(
  484. b'Noise_XK_448_ChaChaPoly_SHA256')
  485. proto.set_as_initiator()
  486. # Setup wrong client key
  487. proto.set_keypair_from_private_bytes(Keypair.STATIC,
  488. wrongclient_key_pair[1])
  489. # but the correct server key
  490. proto.set_keypair_from_public_bytes(Keypair.REMOTE_STATIC,
  491. self.server_key_pair[0])
  492. proto.start_handshake()
  493. # Send first message
  494. message = proto.write_message()
  495. self.assertEqual(len(message), _handshakelens[0])
  496. writer.write(message)
  497. # Get response
  498. respmsg = await reader.readexactly(_handshakelens[1])
  499. proto.read_message(respmsg)
  500. # Send final reply
  501. message = proto.write_message()
  502. writer.write(message)
  503. # Make sure handshake has completed
  504. self.assertTrue(proto.handshake_finished)
  505. with self.assertRaises(ValueError):
  506. await servnf
  507. @async_test
  508. async def test_server(self):
  509. # Test is plumbed:
  510. # (reader, writer) -> servsock ->
  511. # (rdr, wrr) NoiseForward (reader, writer) ->
  512. # servptsock -> (ptsock[0], ptsock[1])
  513. # Path that the server will sit on
  514. servsockpath = os.path.join(self.tempdir, 'servsock')
  515. servarg = _makeunix(servsockpath)
  516. # Path that the server will send pt data to
  517. servptpath = os.path.join(self.tempdir, 'servptsock')
  518. # Setup pt target listener
  519. pttarg = _makeunix(servptpath)
  520. ptsock = []
  521. def ptsockaccept(reader, writer, ptsock=ptsock):
  522. ptsock.append((reader, writer))
  523. # Bind to pt listener
  524. lsock = await listensockstr(pttarg, ptsockaccept)
  525. nfs = []
  526. event = asyncio.Event()
  527. async def runnf(rdr, wrr):
  528. ptpairfun = asyncio.create_task(connectsockstr(pttarg))
  529. a = await NoiseForwarder('resp',
  530. _makefut((rdr, wrr)), lambda x: ptpairfun,
  531. priv_key=self.server_key_pair[1])
  532. nfs.append(a)
  533. event.set()
  534. # Setup server listener
  535. ssock = await listensockstr(servarg, runnf)
  536. # Connect to server
  537. reader, writer = await connectsockstr(servarg)
  538. # Create client
  539. proto = NoiseConnection.from_name(b'Noise_XK_448_ChaChaPoly_SHA256')
  540. proto.set_as_initiator()
  541. # Setup required keys
  542. proto.set_keypair_from_private_bytes(Keypair.STATIC,
  543. self.client_key_pair[1])
  544. proto.set_keypair_from_public_bytes(Keypair.REMOTE_STATIC,
  545. self.server_key_pair[0])
  546. proto.start_handshake()
  547. # Send first message
  548. message = proto.write_message()
  549. self.assertEqual(len(message), _handshakelens[0])
  550. writer.write(message)
  551. # Get response
  552. respmsg = await reader.readexactly(_handshakelens[1])
  553. proto.read_message(respmsg)
  554. # Send final reply
  555. message = proto.write_message()
  556. writer.write(message)
  557. # Make sure handshake has completed
  558. self.assertTrue(proto.handshake_finished)
  559. # generate the keys for lengths
  560. enclenfun, _ = _genciphfun(proto.get_handshake_hash(),
  561. b'toresp')
  562. _, declenfun = _genciphfun(proto.get_handshake_hash(),
  563. b'toinit')
  564. # write a test message
  565. ptmsg = b'this is a test message that should be a little in length'
  566. encmsg = proto.encrypt(ptmsg)
  567. writer.write(enclenfun(encmsg))
  568. writer.write(encmsg)
  569. # XXX - how to sync?
  570. await asyncio.sleep(.1)
  571. ptreader, ptwriter = ptsock[0]
  572. # read the test message
  573. rptmsg = await ptreader.readexactly(len(ptmsg))
  574. self.assertEqual(rptmsg, ptmsg)
  575. # write a different message
  576. ptmsg = os.urandom(2843)
  577. encmsg = proto.encrypt(ptmsg)
  578. writer.write(enclenfun(encmsg))
  579. writer.write(encmsg)
  580. # XXX - how to sync?
  581. await asyncio.sleep(.1)
  582. # read the test message
  583. rptmsg = await ptreader.readexactly(len(ptmsg))
  584. self.assertEqual(rptmsg, ptmsg)
  585. # now try the other way
  586. ptmsg = os.urandom(912)
  587. ptwriter.write(ptmsg)
  588. # find out how much we need to read
  589. encmsg = await reader.readexactly(2 + 16)
  590. tlen = declenfun(encmsg)
  591. # read the rest of the message
  592. rencmsg = await reader.readexactly(tlen - 16)
  593. tmsg = encmsg[2:] + rencmsg
  594. rptmsg = proto.decrypt(tmsg)
  595. self.assertEqual(rptmsg, ptmsg)
  596. # shut down sending
  597. writer.write_eof()
  598. # so pt reader should be shut down
  599. self.assertEqual(b'', await ptreader.read(1))
  600. self.assertTrue(ptreader.at_eof())
  601. # shut down pt
  602. ptwriter.write_eof()
  603. # make sure the enc reader is eof
  604. self.assertEqual(b'', await reader.read(1))
  605. self.assertTrue(reader.at_eof())
  606. await event.wait()
  607. self.assertEqual(nfs[0], [ 'dec', 'enc' ])
  608. @async_test
  609. async def test_serverclient(self):
  610. # plumbing:
  611. #
  612. # ptca -> ptcb NF client clsa -> clsb NF server ptsa -> ptsb
  613. #
  614. ptcsockapair, ptcsockbpair = _asyncsockpair()
  615. ptcareader, ptcawriter = await ptcsockapair
  616. #ptcsockbpair passed directly
  617. clssockapair, clssockbpair = _asyncsockpair()
  618. #both passed directly
  619. ptssockapair, ptssockbpair = _asyncsockpair()
  620. #ptssockapair passed directly
  621. ptsbreader, ptsbwriter = await ptssockbpair
  622. async def validateclientkey(pubkey):
  623. if pubkey != self.client_key_pair[0]:
  624. raise ValueError('invalid key')
  625. return await ptssockapair
  626. clientnf = asyncio.create_task(NoiseForwarder('init',
  627. clssockapair, lambda x: ptcsockbpair,
  628. priv_key=self.client_key_pair[1],
  629. pub_key=self.server_key_pair[0]))
  630. servnf = asyncio.create_task(NoiseForwarder('resp',
  631. clssockbpair, validateclientkey,
  632. priv_key=self.server_key_pair[1]))
  633. # send a message
  634. msga = os.urandom(183)
  635. ptcawriter.write(msga)
  636. # make sure we get the same message
  637. self.assertEqual(msga, await ptsbreader.readexactly(len(msga)))
  638. # send a second message
  639. msga = os.urandom(2834)
  640. ptcawriter.write(msga)
  641. # make sure we get the same message
  642. self.assertEqual(msga, await ptsbreader.readexactly(len(msga)))
  643. # send a message larger than the block size
  644. msga = os.urandom(103958)
  645. ptcawriter.write(msga)
  646. # make sure we get the same message
  647. self.assertEqual(msga, await ptsbreader.readexactly(len(msga)))
  648. # send a message the other direction
  649. msga = os.urandom(103958)
  650. ptsbwriter.write(msga)
  651. # make sure we get the same message
  652. self.assertEqual(msga, await ptcareader.readexactly(len(msga)))
  653. # close down the pt writers, the rest should follow
  654. ptsbwriter.write_eof()
  655. ptcawriter.write_eof()
  656. # make sure they are closed, and there is no more data
  657. self.assertEqual(b'', await ptsbreader.read(1))
  658. self.assertTrue(ptsbreader.at_eof())
  659. self.assertEqual(b'', await ptcareader.read(1))
  660. self.assertTrue(ptcareader.at_eof())
  661. self.assertEqual([ 'dec', 'enc' ], await clientnf)
  662. self.assertEqual([ 'dec', 'enc' ], await servnf)