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.
 
 

1203 lines
32 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(100)
  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. # Make sure any additions are reflected by tests in test_parsesockstr
  62. _allowedparameters = {
  63. 'unix': {
  64. 'path': str,
  65. },
  66. 'tcp': {
  67. 'host': str,
  68. 'port': int,
  69. },
  70. }
  71. def parsesockstr(sockstr):
  72. '''Parse a socket string to its parts.
  73. The format of sockstr is: 'proto:param=value[,param2=value2]'.
  74. If the proto has a default parameter, the value can be used
  75. directly, like: 'proto:value'. This is only allowed when the
  76. value can unambiguously be determined not to be a param. If
  77. there needs to be an equals '=', then you MUST use the extended
  78. version.
  79. The characters that define 'param' must be all lower case ascii
  80. characters and may contain an underscore. The first character
  81. must not be an underscore.
  82. Supported protocols:
  83. unix:
  84. Default parameter is path.
  85. The path parameter specifies the path to the
  86. unix domain socket. The path MUST start w/ a
  87. slash if it is used as a default parameter.
  88. tcp:
  89. Default parameter is host[:port].
  90. The host parameter specifies the host, and the
  91. port parameter specifies the port of the
  92. connection.
  93. '''
  94. proto, rem = sockstr.split(':', 1)
  95. if '=' not in rem:
  96. if proto == 'unix' and rem[0] != '/':
  97. raise ValueError('bare path MUST start w/ a slash (/).')
  98. if proto == 'unix':
  99. args = { 'path': rem }
  100. else:
  101. args = dict(i.split('=', 1) for i in rem.split(','))
  102. try:
  103. allowed = _allowedparameters[proto]
  104. except KeyError:
  105. raise ValueError('unsupported proto: %s' % repr(proto))
  106. extrakeys = args.keys() - allowed.keys()
  107. if extrakeys:
  108. raise ValueError('keys for proto %s not allowed: %s' % (repr(proto), extrakeys))
  109. for i in args:
  110. args[i] = allowed[i](args[i])
  111. return proto, args
  112. async def connectsockstr(sockstr):
  113. '''Wrapper for asyncio.open_*_connection.'''
  114. proto, args = parsesockstr(sockstr)
  115. if proto == 'unix':
  116. fun = asyncio.open_unix_connection
  117. elif proto == 'tcp':
  118. fun = asyncio.open_connection
  119. reader, writer = await fun(**args)
  120. return reader, writer
  121. async def listensockstr(sockstr, cb):
  122. '''Wrapper for asyncio.start_x_server.
  123. For the format of sockstr, please see parsesockstr.
  124. The cb parameter is passed to asyncio's start_server or related
  125. calls. Per those docs, the cb parameter is calls or scheduled
  126. as a task when a client establishes a connection. It is called
  127. with two arguments, the reader and writer streams. For more
  128. information, see: https://docs.python.org/3/library/asyncio-stream.html#asyncio.start_server
  129. '''
  130. proto, args = parsesockstr(sockstr)
  131. if proto == 'unix':
  132. fun = asyncio.start_unix_server
  133. elif proto == 'tcp':
  134. fun = asyncio.start_server
  135. return await fun(cb, **args)
  136. # !!python makemessagelengths.py
  137. _handshakelens = \
  138. [72, 72, 88]
  139. def _genciphfun(hash, ad):
  140. hkdf = HKDF(algorithm=hashes.SHA256(), length=32,
  141. salt=b'asdoifjsldkjdsf', info=ad, backend=_backend)
  142. key = hkdf.derive(hash)
  143. cipher = Cipher(algorithms.AES(key), modes.ECB(),
  144. backend=_backend)
  145. enctor = cipher.encryptor()
  146. def encfun(data):
  147. # Returns the two bytes for length
  148. val = len(data)
  149. encbytes = enctor.update(data[:16])
  150. mask = int.from_bytes(encbytes[:2], byteorder='big') & 0xff
  151. return (val ^ mask).to_bytes(length=2, byteorder='big')
  152. def decfun(data):
  153. # takes off the data and returns the total
  154. # length
  155. val = int.from_bytes(data[:2], byteorder='big')
  156. encbytes = enctor.update(data[2:2 + 16])
  157. mask = int.from_bytes(encbytes[:2], byteorder='big') & 0xff
  158. return val ^ mask
  159. return encfun, decfun
  160. async def NoiseForwarder(mode, encrdrwrr, ptpairfun, priv_key, pub_key=None):
  161. '''A function that forwards data between the plain text pair of
  162. streams to the encrypted session.
  163. The mode paramater must be one of 'init' or 'resp' for initiator
  164. and responder.
  165. The encrdrwrr is an await object that will return a tunle of the
  166. reader and writer streams for the encrypted side of the
  167. connection.
  168. The ptpairfun parameter is a function that will be passed the
  169. public key bytes for the remote client. This can be used to
  170. both validate that the correct client is connecting, and to
  171. pass back the correct plain text reader/writer objects that
  172. match the provided static key. The function must be an async
  173. function.
  174. In the case of the initiator, pub_key must be provided and will
  175. be used to authenticate the responder side of the connection.
  176. The priv_key parameter is used to authenticate this side of the
  177. session.
  178. Both priv_key and pub_key parameters must be 56 bytes. For example,
  179. the pair that is returned by genkeypair.
  180. '''
  181. # Send a protocol version so that in the future we can change how
  182. # we interface, and possibly be able to send control messages,
  183. # allow the client to pass some misc data to the callback, or to
  184. # allow a reverse tunnel, were the client talks to the server,
  185. # and waits for the server to "connect" to the client w/ a
  186. # connection, e.g. reverse tunnel out behind a nat to allow
  187. # incoming connections.
  188. protocol_version = 0
  189. rdr, wrr = await encrdrwrr
  190. proto = NoiseConnection.from_name(b'Noise_XK_448_ChaChaPoly_SHA256')
  191. proto.set_keypair_from_private_bytes(Keypair.STATIC, priv_key)
  192. if pub_key is not None:
  193. proto.set_keypair_from_public_bytes(Keypair.REMOTE_STATIC,
  194. pub_key)
  195. if mode == 'resp':
  196. proto.set_as_responder()
  197. proto.start_handshake()
  198. proto.read_message(await rdr.readexactly(_handshakelens[0]))
  199. wrr.write(proto.write_message())
  200. proto.read_message(await rdr.readexactly(_handshakelens[2]))
  201. elif mode == 'init':
  202. proto.set_as_initiator()
  203. proto.start_handshake()
  204. wrr.write(proto.write_message())
  205. proto.read_message(await rdr.readexactly(_handshakelens[1]))
  206. wrr.write(proto.write_message())
  207. if not proto.handshake_finished: # pragma: no cover
  208. raise RuntimeError('failed to finish handshake')
  209. try:
  210. reader, writer = await ptpairfun(getattr(proto.get_keypair(
  211. Keypair.REMOTE_STATIC), 'public_bytes', None))
  212. except:
  213. wrr.close()
  214. raise
  215. # generate the keys for lengths
  216. # XXX - get_handshake_hash is probably not the best option, but
  217. # this is only to obscure lengths, it is not required to be secure
  218. # as the underlying NoiseProtocol securely validates everything.
  219. # It is marginally useful as writing patterns likely expose the
  220. # true length. Adding padding could marginally help w/ this.
  221. if mode == 'resp':
  222. _, declenfun = _genciphfun(proto.get_handshake_hash(), b'toresp')
  223. enclenfun, _ = _genciphfun(proto.get_handshake_hash(), b'toinit')
  224. elif mode == 'init':
  225. enclenfun, _ = _genciphfun(proto.get_handshake_hash(), b'toresp')
  226. _, declenfun = _genciphfun(proto.get_handshake_hash(), b'toinit')
  227. # protocol negotiation
  228. # send first, then wait for the response
  229. pvmsg = protocol_version.to_bytes(1, byteorder='big')
  230. encmsg = proto.encrypt(pvmsg)
  231. wrr.write(enclenfun(encmsg))
  232. wrr.write(encmsg)
  233. # get the protocol version
  234. msg = await rdr.readexactly(2 + 16)
  235. tlen = declenfun(msg)
  236. rmsg = await rdr.readexactly(tlen - 16)
  237. tmsg = msg[2:] + rmsg
  238. rpv = proto.decrypt(tmsg)
  239. async def decses():
  240. try:
  241. while True:
  242. try:
  243. msg = await rdr.readexactly(2 + 16)
  244. except asyncio.streams.IncompleteReadError:
  245. if rdr.at_eof():
  246. return 'dec'
  247. tlen = declenfun(msg)
  248. rmsg = await rdr.readexactly(tlen - 16)
  249. tmsg = msg[2:] + rmsg
  250. writer.write(proto.decrypt(tmsg))
  251. await writer.drain()
  252. #except:
  253. # import traceback
  254. # traceback.print_exc()
  255. # raise
  256. finally:
  257. try:
  258. writer.write_eof()
  259. except OSError as e:
  260. if e.errno != 57:
  261. raise
  262. async def encses():
  263. try:
  264. while True:
  265. # largest message
  266. ptmsg = await reader.read(65535 - 16)
  267. if not ptmsg:
  268. # eof
  269. return 'enc'
  270. encmsg = proto.encrypt(ptmsg)
  271. wrr.write(enclenfun(encmsg))
  272. wrr.write(encmsg)
  273. await wrr.drain()
  274. #except:
  275. # import traceback
  276. # traceback.print_exc()
  277. # raise
  278. finally:
  279. wrr.write_eof()
  280. res = await asyncio.gather(decses(), encses())
  281. await wrr.drain() # not sure if needed
  282. wrr.close()
  283. await writer.drain() # not sure if needed
  284. writer.close()
  285. return res
  286. # https://stackoverflow.com/questions/23033939/how-to-test-python-3-4-asyncio-code
  287. # Slightly modified to timeout and to print trace back when canceled.
  288. # This makes it easier to figure out what "froze".
  289. def async_test(f):
  290. def wrapper(*args, **kwargs):
  291. async def tbcapture():
  292. try:
  293. return await f(*args, **kwargs)
  294. except asyncio.CancelledError as e:
  295. # if we are going to be cancelled, print out a tb
  296. import traceback
  297. traceback.print_exc()
  298. raise
  299. loop = asyncio.get_event_loop()
  300. # timeout after 4 seconds
  301. loop.run_until_complete(asyncio.wait_for(tbcapture(), 4))
  302. return wrapper
  303. class Tests_misc(unittest.TestCase):
  304. def setUp(self):
  305. # setup temporary directory
  306. d = os.path.realpath(tempfile.mkdtemp())
  307. self.basetempdir = d
  308. self.tempdir = os.path.join(d, 'subdir')
  309. os.mkdir(self.tempdir)
  310. os.chdir(self.tempdir)
  311. def tearDown(self):
  312. #print('td:', time.time())
  313. shutil.rmtree(self.basetempdir)
  314. self.tempdir = None
  315. def test_parsesockstr_bad(self):
  316. badstrs = [
  317. 'unix:ff',
  318. 'randomnocolon',
  319. 'unix:somethingelse=bogus',
  320. 'tcp:port=bogus',
  321. ]
  322. for i in badstrs:
  323. with self.assertRaises(ValueError,
  324. msg='Should have failed processing: %s' % repr(i)):
  325. parsesockstr(i)
  326. def test_parsesockstr(self):
  327. results = {
  328. # Not all of these are valid when passed to a *sockstr
  329. # function
  330. 'unix:/apath': ('unix', { 'path': '/apath' }),
  331. 'unix:path=apath': ('unix', { 'path': 'apath' }),
  332. 'tcp:host=apath': ('tcp', { 'host': 'apath' }),
  333. 'tcp:host=apath,port=5': ('tcp', { 'host': 'apath',
  334. 'port': 5 }),
  335. }
  336. for s, r in results.items():
  337. self.assertEqual(parsesockstr(s), r)
  338. @async_test
  339. async def test_listensockstr_bad(self):
  340. with self.assertRaises(ValueError):
  341. ls = await listensockstr('bogus:some=arg', None)
  342. with self.assertRaises(ValueError):
  343. ls = await connectsockstr('bogus:some=arg')
  344. @async_test
  345. async def test_listenconnectsockstr(self):
  346. msgsent = b'this is a test message'
  347. msgrcv = b'testing message for receive'
  348. # That when a connection is received and receives and sends
  349. async def servconfhandle(rdr, wrr):
  350. msg = await rdr.readexactly(len(msgsent))
  351. self.assertEqual(msg, msgsent)
  352. #print(repr(wrr.get_extra_info('sockname')))
  353. wrr.write(msgrcv)
  354. await wrr.drain()
  355. wrr.close()
  356. return True
  357. # Test listensockstr
  358. for sstr, confun in [
  359. ('unix:path=ff', lambda: asyncio.open_unix_connection(path='ff')),
  360. ('tcp:port=9384', lambda: asyncio.open_connection(port=9384))
  361. ]:
  362. # that listensockstr will bind to the correct path, can call cb
  363. ls = await listensockstr(sstr, servconfhandle)
  364. # that we open a connection to the path
  365. rdr, wrr = await confun()
  366. # and send a message
  367. wrr.write(msgsent)
  368. # and receive the message
  369. rcv = await asyncio.wait_for(rdr.readexactly(len(msgrcv)), .5)
  370. self.assertEqual(rcv, msgrcv)
  371. wrr.close()
  372. # Now test that connectsockstr works similarly.
  373. rdr, wrr = await connectsockstr(sstr)
  374. # and send a message
  375. wrr.write(msgsent)
  376. # and receive the message
  377. rcv = await asyncio.wait_for(rdr.readexactly(len(msgrcv)), .5)
  378. self.assertEqual(rcv, msgrcv)
  379. wrr.close()
  380. ls.close()
  381. await ls.wait_closed()
  382. def test_genciphfun(self):
  383. enc, dec = _genciphfun(b'0' * 32, b'foobar')
  384. msg = b'this is a bunch of data'
  385. tb = enc(msg)
  386. self.assertEqual(len(msg), dec(tb + msg))
  387. for i in [ 20, 1384, 64000, 23839, 65535 ]:
  388. msg = os.urandom(i)
  389. self.assertEqual(len(msg), dec(enc(msg) + msg))
  390. def cmd_client(args):
  391. privkey = loadprivkeyraw(args.clientkey)
  392. pubkey = loadpubkeyraw(args.servkey)
  393. async def runnf(rdr, wrr):
  394. encpair = asyncio.create_task(connectsockstr(args.clienttarget))
  395. a = await NoiseForwarder('init',
  396. encpair, lambda x: _makefut((rdr, wrr)),
  397. priv_key=privkey, pub_key=pubkey)
  398. # Setup client listener
  399. ssock = listensockstr(args.clientlisten, runnf)
  400. loop = asyncio.get_event_loop()
  401. obj = loop.run_until_complete(ssock)
  402. loop.run_until_complete(obj.serve_forever())
  403. def cmd_server(args):
  404. privkey = loadprivkeyraw(args.servkey)
  405. pubkeys = [ loadpubkeyraw(x) for x in args.clientkey ]
  406. async def runnf(rdr, wrr):
  407. async def checkclientfun(clientkey):
  408. if clientkey not in pubkeys:
  409. raise RuntimeError('invalid key provided')
  410. return await connectsockstr(args.servtarget)
  411. a = await NoiseForwarder('resp', _makefut((rdr, wrr)),
  412. checkclientfun, priv_key=privkey)
  413. # Setup server listener
  414. ssock = listensockstr(args.servlisten, runnf)
  415. loop = asyncio.get_event_loop()
  416. obj = loop.run_until_complete(ssock)
  417. loop.run_until_complete(obj.serve_forever())
  418. def cmd_genkey(args):
  419. keypair = genkeypair()
  420. key = x448.X448PrivateKey.generate()
  421. # public key part
  422. enc = serialization.Encoding.Raw
  423. pubformat = serialization.PublicFormat.Raw
  424. pub = key.public_key().public_bytes(encoding=enc, format=pubformat)
  425. try:
  426. fname = args.fname + '.pub'
  427. with open(fname, 'x', encoding='ascii') as fp:
  428. print('ntun-x448', base64.urlsafe_b64encode(pub).decode('ascii'), file=fp)
  429. except FileExistsError:
  430. print('failed to create %s, file exists.' % fname, file=sys.stderr)
  431. sys.exit(2)
  432. enc = serialization.Encoding.PEM
  433. format = serialization.PrivateFormat.PKCS8
  434. encalgo = serialization.NoEncryption()
  435. with open(args.fname, 'x', encoding='ascii') as fp:
  436. fp.write(key.private_bytes(encoding=enc, format=format, encryption_algorithm=encalgo).decode('ascii'))
  437. def main():
  438. parser = argparse.ArgumentParser()
  439. subparsers = parser.add_subparsers(title='subcommands', description='valid subcommands', help='additional help')
  440. parser_gk = subparsers.add_parser('genkey', help='generate keys')
  441. parser_gk.add_argument('fname', type=str, help='file name for the key')
  442. parser_gk.set_defaults(func=cmd_genkey)
  443. parser_serv = subparsers.add_parser('server', help='run a server')
  444. parser_serv.add_argument('--clientkey', '-c', action='append', type=str, help='file of authorized client keys, or a .pub file')
  445. parser_serv.add_argument('servkey', type=str, help='file name for the server key')
  446. parser_serv.add_argument('servlisten', type=str, help='Connection that the server listens on')
  447. parser_serv.add_argument('servtarget', type=str, help='Connection that the server connects to')
  448. parser_serv.set_defaults(func=cmd_server)
  449. parser_client = subparsers.add_parser('client', help='run a client')
  450. parser_client.add_argument('clientkey', type=str, help='file name for the client private key')
  451. parser_client.add_argument('servkey', type=str, help='file name for the server public key')
  452. parser_client.add_argument('clientlisten', type=str, help='Connection that the client listens on')
  453. parser_client.add_argument('clienttarget', type=str, help='Connection that the client connects to')
  454. parser_client.set_defaults(func=cmd_client)
  455. args = parser.parse_args()
  456. try:
  457. fun = args.func
  458. except AttributeError:
  459. parser.print_usage()
  460. sys.exit(5)
  461. fun(args)
  462. if __name__ == '__main__': # pragma: no cover
  463. main()
  464. def _asyncsockpair():
  465. '''Create a pair of sockets that are bound to each other.
  466. The function will return a tuple of two coroutine's, that
  467. each, when await'ed upon, will return the reader/writer pair.'''
  468. socka, sockb = socket.socketpair()
  469. return asyncio.open_connection(sock=socka), \
  470. asyncio.open_connection(sock=sockb)
  471. async def _awaitfile(fname):
  472. while not os.path.exists(fname):
  473. await asyncio.sleep(.01)
  474. return True
  475. class TestMain(unittest.TestCase):
  476. def setUp(self):
  477. # setup temporary directory
  478. d = os.path.realpath(tempfile.mkdtemp())
  479. self.basetempdir = d
  480. self.tempdir = os.path.join(d, 'subdir')
  481. os.mkdir(self.tempdir)
  482. # Generate key pairs
  483. self.server_key_pair = genkeypair()
  484. self.client_key_pair = genkeypair()
  485. os.chdir(self.tempdir)
  486. def tearDown(self):
  487. #print('td:', time.time())
  488. shutil.rmtree(self.basetempdir)
  489. self.tempdir = None
  490. @async_test
  491. async def test_noargs(self):
  492. proc = await self.run_with_args()
  493. await proc.wait()
  494. # XXX - not checking error message
  495. # And that it exited w/ the correct code
  496. self.assertEqual(proc.returncode, 5)
  497. def run_with_args(self, *args, pipes=True):
  498. kwargs = {}
  499. if pipes:
  500. kwargs.update(dict(
  501. stdout=asyncio.subprocess.PIPE,
  502. stderr=asyncio.subprocess.PIPE))
  503. return asyncio.create_subprocess_exec(sys.executable,
  504. # XXX - figure out how to add coverage data on these runs
  505. #'-m', 'coverage', 'run', '-p',
  506. __file__, *args, **kwargs)
  507. async def genkey(self, name):
  508. proc = await self.run_with_args('genkey', name, pipes=False)
  509. await proc.wait()
  510. self.assertEqual(proc.returncode, 0)
  511. @async_test
  512. async def test_loadpubkey(self):
  513. keypath = os.path.join(self.tempdir, 'loadpubkeytest')
  514. await self.genkey(keypath)
  515. privkey = loadprivkey(keypath)
  516. enc = serialization.Encoding.Raw
  517. pubformat = serialization.PublicFormat.Raw
  518. pubkeybytes = privkey.public_key().public_bytes(encoding=enc, format=pubformat)
  519. pubkey = loadpubkeyraw(keypath + '.pub')
  520. self.assertEqual(pubkeybytes, pubkey)
  521. privrawkey = loadprivkeyraw(keypath)
  522. enc = serialization.Encoding.Raw
  523. privformat = serialization.PrivateFormat.Raw
  524. encalgo = serialization.NoEncryption()
  525. rprivrawkey = privkey.private_bytes(encoding=enc, format=privformat, encryption_algorithm=encalgo)
  526. self.assertEqual(rprivrawkey, privrawkey)
  527. @async_test
  528. async def test_clientkeymismatch(self):
  529. # make sure that if there's a client key mismatch, we
  530. # don't connect
  531. # Generate necessar keys
  532. servkeypath = os.path.join(self.tempdir, 'server_key')
  533. await self.genkey(servkeypath)
  534. clientkeypath = os.path.join(self.tempdir, 'client_key')
  535. await self.genkey(clientkeypath)
  536. badclientkeypath = os.path.join(self.tempdir, 'badclient_key')
  537. await self.genkey(badclientkeypath)
  538. # forwards connectsion to this socket (created by client)
  539. ptclientpath = os.path.join(self.tempdir, 'incclient.sock')
  540. ptclientstr = _makeunix(ptclientpath)
  541. # this is the socket server listen to
  542. incservpath = os.path.join(self.tempdir, 'incserv.sock')
  543. incservstr = _makeunix(incservpath)
  544. # to this socket, opened by server
  545. servtargpath = os.path.join(self.tempdir, 'servtarget.sock')
  546. servtargstr = _makeunix(servtargpath)
  547. # Setup server target listener
  548. ptsockevent = asyncio.Event()
  549. # Bind to pt listener
  550. lsock = await listensockstr(servtargstr, None)
  551. # Startup the server
  552. server = await self.run_with_args('server',
  553. '-c', clientkeypath + '.pub',
  554. servkeypath, incservstr, servtargstr)
  555. # Startup the client with the "bad" key
  556. client = await self.run_with_args('client',
  557. badclientkeypath, servkeypath + '.pub', ptclientstr, incservstr)
  558. # wait for server target to be created
  559. await _awaitfile(servtargpath)
  560. # wait for server to start
  561. await _awaitfile(incservpath)
  562. # wait for client to start
  563. await _awaitfile(ptclientpath)
  564. # Connect to the client
  565. reader, writer = await connectsockstr(ptclientstr)
  566. # XXX - this might not be the best test.
  567. with self.assertRaises(asyncio.futures.TimeoutError):
  568. # make sure that we don't get the conenction
  569. await asyncio.wait_for(ptsockevent.wait(), .5)
  570. writer.close()
  571. # Make sure that when the server is terminated
  572. server.terminate()
  573. # that it's stderr
  574. stdout, stderr = await server.communicate()
  575. #print('s:', repr((stdout, stderr)))
  576. # doesn't have an exceptions never retrieved
  577. # even the example echo server has this same leak
  578. #self.assertNotIn(b'Task exception was never retrieved', stderr)
  579. lsock.close()
  580. await lsock.wait_closed()
  581. @async_test
  582. async def test_end2end(self):
  583. # Generate necessar keys
  584. servkeypath = os.path.join(self.tempdir, 'server_key')
  585. await self.genkey(servkeypath)
  586. clientkeypath = os.path.join(self.tempdir, 'client_key')
  587. await self.genkey(clientkeypath)
  588. # forwards connectsion to this socket (created by client)
  589. ptclientpath = os.path.join(self.tempdir, 'incclient.sock')
  590. ptclientstr = _makeunix(ptclientpath)
  591. # this is the socket server listen to
  592. incservpath = os.path.join(self.tempdir, 'incserv.sock')
  593. incservstr = _makeunix(incservpath)
  594. # to this socket, opened by server
  595. servtargpath = os.path.join(self.tempdir, 'servtarget.sock')
  596. servtargstr = _makeunix(servtargpath)
  597. # Setup server target listener
  598. ptsock = []
  599. ptsockevent = asyncio.Event()
  600. def ptsockaccept(reader, writer, ptsock=ptsock):
  601. ptsock.append((reader, writer))
  602. ptsockevent.set()
  603. # Bind to pt listener
  604. lsock = await listensockstr(servtargstr, ptsockaccept)
  605. # Startup the server
  606. server = await self.run_with_args('server',
  607. '-c', clientkeypath + '.pub',
  608. servkeypath, incservstr, servtargstr,
  609. pipes=False)
  610. # Startup the client
  611. client = await self.run_with_args('client',
  612. clientkeypath, servkeypath + '.pub', ptclientstr, incservstr,
  613. pipes=False)
  614. # wait for server target to be created
  615. await _awaitfile(servtargpath)
  616. # wait for server to start
  617. await _awaitfile(incservpath)
  618. # wait for client to start
  619. await _awaitfile(ptclientpath)
  620. # Connect to the client
  621. reader, writer = await connectsockstr(ptclientstr)
  622. # send a message
  623. ptmsg = b'this is a message for testing'
  624. writer.write(ptmsg)
  625. # make sure that we got the conenction
  626. await ptsockevent.wait()
  627. # get the connection
  628. endrdr, endwrr = ptsock[0]
  629. # make sure we can read back what we sent
  630. self.assertEqual(ptmsg, await endrdr.readexactly(len(ptmsg)))
  631. # test some additional messages
  632. for i in [ 129, 1287, 28792, 129872 ]:
  633. # in on direction
  634. msg = os.urandom(i)
  635. writer.write(msg)
  636. self.assertEqual(msg, await endrdr.readexactly(len(msg)))
  637. # and the other
  638. endwrr.write(msg)
  639. self.assertEqual(msg, await reader.readexactly(len(msg)))
  640. writer.close()
  641. endwrr.close()
  642. lsock.close()
  643. await lsock.wait_closed()
  644. @async_test
  645. async def test_genkey(self):
  646. # that it can generate a key
  647. proc = await self.run_with_args('genkey', 'somefile')
  648. await proc.wait()
  649. #print(await proc.communicate())
  650. self.assertEqual(proc.returncode, 0)
  651. with open('somefile.pub', encoding='ascii') as fp:
  652. lines = fp.readlines()
  653. self.assertEqual(len(lines), 1)
  654. keytype, keyvalue = lines[0].split()
  655. self.assertEqual(keytype, 'ntun-x448')
  656. key = x448.X448PublicKey.from_public_bytes(base64.urlsafe_b64decode(keyvalue))
  657. key = loadprivkey('somefile')
  658. self.assertIsInstance(key, x448.X448PrivateKey)
  659. # that a second call fails
  660. proc = await self.run_with_args('genkey', 'somefile')
  661. await proc.wait()
  662. stdoutdata, stderrdata = await proc.communicate()
  663. self.assertFalse(stdoutdata)
  664. self.assertEqual(b'failed to create somefile.pub, file exists.\n', stderrdata)
  665. # And that it exited w/ the correct code
  666. self.assertEqual(proc.returncode, 2)
  667. class TestNoiseFowarder(unittest.TestCase):
  668. def setUp(self):
  669. # setup temporary directory
  670. d = os.path.realpath(tempfile.mkdtemp())
  671. self.basetempdir = d
  672. self.tempdir = os.path.join(d, 'subdir')
  673. os.mkdir(self.tempdir)
  674. # Generate key pairs
  675. self.server_key_pair = genkeypair()
  676. self.client_key_pair = genkeypair()
  677. def tearDown(self):
  678. shutil.rmtree(self.basetempdir)
  679. self.tempdir = None
  680. @async_test
  681. async def test_clientkeymissmatch(self):
  682. # generate a key that is incorrect
  683. wrongclient_key_pair = genkeypair()
  684. # the secure socket
  685. clssockapair, clssockbpair = _asyncsockpair()
  686. reader, writer = await clssockapair
  687. async def wrongkey(v):
  688. raise ValueError('no key matches')
  689. # create the server
  690. servnf = asyncio.create_task(NoiseForwarder('resp',
  691. clssockbpair, wrongkey,
  692. priv_key=self.server_key_pair[1]))
  693. # Create client
  694. proto = NoiseConnection.from_name(
  695. b'Noise_XK_448_ChaChaPoly_SHA256')
  696. proto.set_as_initiator()
  697. # Setup wrong client key
  698. proto.set_keypair_from_private_bytes(Keypair.STATIC,
  699. wrongclient_key_pair[1])
  700. # but the correct server key
  701. proto.set_keypair_from_public_bytes(Keypair.REMOTE_STATIC,
  702. self.server_key_pair[0])
  703. proto.start_handshake()
  704. # Send first message
  705. message = proto.write_message()
  706. self.assertEqual(len(message), _handshakelens[0])
  707. writer.write(message)
  708. # Get response
  709. respmsg = await reader.readexactly(_handshakelens[1])
  710. proto.read_message(respmsg)
  711. # Send final reply
  712. message = proto.write_message()
  713. writer.write(message)
  714. # Make sure handshake has completed
  715. self.assertTrue(proto.handshake_finished)
  716. with self.assertRaises(ValueError):
  717. await servnf
  718. writer.close()
  719. @async_test
  720. async def test_server(self):
  721. # Test is plumbed:
  722. # (reader, writer) -> servsock ->
  723. # (rdr, wrr) NoiseForward (reader, writer) ->
  724. # servptsock -> (ptsock[0], ptsock[1])
  725. # Path that the server will sit on
  726. servsockpath = os.path.join(self.tempdir, 'servsock')
  727. servarg = _makeunix(servsockpath)
  728. # Path that the server will send pt data to
  729. servptpath = os.path.join(self.tempdir, 'servptsock')
  730. # Setup pt target listener
  731. pttarg = _makeunix(servptpath)
  732. ptsock = []
  733. ptsockevent = asyncio.Event()
  734. def ptsockaccept(reader, writer, ptsock=ptsock):
  735. ptsock.append((reader, writer))
  736. ptsockevent.set()
  737. # Bind to pt listener
  738. lsock = await listensockstr(pttarg, ptsockaccept)
  739. nfs = []
  740. event = asyncio.Event()
  741. async def runnf(rdr, wrr):
  742. ptpairfun = asyncio.create_task(connectsockstr(pttarg))
  743. a = await NoiseForwarder('resp',
  744. _makefut((rdr, wrr)), lambda x: ptpairfun,
  745. priv_key=self.server_key_pair[1])
  746. nfs.append(a)
  747. event.set()
  748. # Setup server listener
  749. ssock = await listensockstr(servarg, runnf)
  750. # Connect to server
  751. reader, writer = await connectsockstr(servarg)
  752. # Create client
  753. proto = NoiseConnection.from_name(b'Noise_XK_448_ChaChaPoly_SHA256')
  754. proto.set_as_initiator()
  755. # Setup required keys
  756. proto.set_keypair_from_private_bytes(Keypair.STATIC,
  757. self.client_key_pair[1])
  758. proto.set_keypair_from_public_bytes(Keypair.REMOTE_STATIC,
  759. self.server_key_pair[0])
  760. proto.start_handshake()
  761. # Send first message
  762. message = proto.write_message()
  763. self.assertEqual(len(message), _handshakelens[0])
  764. writer.write(message)
  765. # Get response
  766. respmsg = await reader.readexactly(_handshakelens[1])
  767. proto.read_message(respmsg)
  768. # Send final reply
  769. message = proto.write_message()
  770. writer.write(message)
  771. # Make sure handshake has completed
  772. self.assertTrue(proto.handshake_finished)
  773. # generate the keys for lengths
  774. enclenfun, _ = _genciphfun(proto.get_handshake_hash(),
  775. b'toresp')
  776. _, declenfun = _genciphfun(proto.get_handshake_hash(),
  777. b'toinit')
  778. pversion = 0
  779. # Send the protocol version string first
  780. encmsg = proto.encrypt(pversion.to_bytes(1, byteorder='big'))
  781. writer.write(enclenfun(encmsg))
  782. writer.write(encmsg)
  783. # Read the peer's protocol version
  784. # find out how much we need to read
  785. encmsg = await reader.readexactly(2 + 16)
  786. tlen = declenfun(encmsg)
  787. # read the rest of the message
  788. rencmsg = await reader.readexactly(tlen - 16)
  789. tmsg = encmsg[2:] + rencmsg
  790. rptmsg = proto.decrypt(tmsg)
  791. self.assertEqual(int.from_bytes(rptmsg, byteorder='big'), pversion)
  792. # write a test message
  793. ptmsg = b'this is a test message that should be a little in length'
  794. encmsg = proto.encrypt(ptmsg)
  795. writer.write(enclenfun(encmsg))
  796. writer.write(encmsg)
  797. # wait for the connection to arrive
  798. await ptsockevent.wait()
  799. ptreader, ptwriter = ptsock[0]
  800. # read the test message
  801. rptmsg = await ptreader.readexactly(len(ptmsg))
  802. self.assertEqual(rptmsg, ptmsg)
  803. # write a different message
  804. ptmsg = os.urandom(2843)
  805. encmsg = proto.encrypt(ptmsg)
  806. writer.write(enclenfun(encmsg))
  807. writer.write(encmsg)
  808. # read the test message
  809. rptmsg = await ptreader.readexactly(len(ptmsg))
  810. self.assertEqual(rptmsg, ptmsg)
  811. # now try the other way
  812. ptmsg = os.urandom(912)
  813. ptwriter.write(ptmsg)
  814. # find out how much we need to read
  815. encmsg = await reader.readexactly(2 + 16)
  816. tlen = declenfun(encmsg)
  817. # read the rest of the message
  818. rencmsg = await reader.readexactly(tlen - 16)
  819. tmsg = encmsg[2:] + rencmsg
  820. rptmsg = proto.decrypt(tmsg)
  821. self.assertEqual(rptmsg, ptmsg)
  822. # shut down sending
  823. writer.write_eof()
  824. # so pt reader should be shut down
  825. self.assertEqual(b'', await ptreader.read(1))
  826. self.assertTrue(ptreader.at_eof())
  827. # shut down pt
  828. ptwriter.write_eof()
  829. # make sure the enc reader is eof
  830. self.assertEqual(b'', await reader.read(1))
  831. self.assertTrue(reader.at_eof())
  832. await event.wait()
  833. self.assertEqual(nfs[0], [ 'dec', 'enc' ])
  834. writer.close()
  835. ptwriter.close()
  836. lsock.close()
  837. ssock.close()
  838. await lsock.wait_closed()
  839. await ssock.wait_closed()
  840. @async_test
  841. async def test_serverclient(self):
  842. # plumbing:
  843. #
  844. # ptca -> ptcb NF client clsa -> clsb NF server ptsa -> ptsb
  845. #
  846. ptcsockapair, ptcsockbpair = _asyncsockpair()
  847. ptcareader, ptcawriter = await ptcsockapair
  848. #ptcsockbpair passed directly
  849. clssockapair, clssockbpair = _asyncsockpair()
  850. #both passed directly
  851. ptssockapair, ptssockbpair = _asyncsockpair()
  852. #ptssockapair passed directly
  853. ptsbreader, ptsbwriter = await ptssockbpair
  854. async def validateclientkey(pubkey):
  855. self.assertEqual(pubkey, self.client_key_pair[0])
  856. return await ptssockapair
  857. clientnf = asyncio.create_task(NoiseForwarder('init',
  858. clssockapair, lambda x: ptcsockbpair,
  859. priv_key=self.client_key_pair[1],
  860. pub_key=self.server_key_pair[0]))
  861. servnf = asyncio.create_task(NoiseForwarder('resp',
  862. clssockbpair, validateclientkey,
  863. priv_key=self.server_key_pair[1]))
  864. # send a message
  865. msga = os.urandom(183)
  866. ptcawriter.write(msga)
  867. # make sure we get the same message
  868. self.assertEqual(msga, await ptsbreader.readexactly(len(msga)))
  869. # send a second message
  870. msga = os.urandom(2834)
  871. ptcawriter.write(msga)
  872. # make sure we get the same message
  873. self.assertEqual(msga, await ptsbreader.readexactly(len(msga)))
  874. # send a message larger than the block size
  875. msga = os.urandom(103958)
  876. ptcawriter.write(msga)
  877. # make sure we get the same message
  878. self.assertEqual(msga, await ptsbreader.readexactly(len(msga)))
  879. # send a message the other direction
  880. msga = os.urandom(103958)
  881. ptsbwriter.write(msga)
  882. # make sure we get the same message
  883. self.assertEqual(msga, await ptcareader.readexactly(len(msga)))
  884. # close down the pt writers, the rest should follow
  885. ptsbwriter.write_eof()
  886. ptcawriter.write_eof()
  887. # make sure they are closed, and there is no more data
  888. self.assertEqual(b'', await ptsbreader.read(1))
  889. self.assertTrue(ptsbreader.at_eof())
  890. self.assertEqual(b'', await ptcareader.read(1))
  891. self.assertTrue(ptcareader.at_eof())
  892. self.assertEqual([ 'dec', 'enc' ], await clientnf)
  893. self.assertEqual([ 'dec', 'enc' ], await servnf)
  894. await ptsbwriter.drain()
  895. await ptcawriter.drain()
  896. ptsbwriter.close()
  897. ptcawriter.close()