From 523132885f6af8a108b7d7f8eed689af52089acc Mon Sep 17 00:00:00 2001
From: John-Mark Gurney <jmg@funkthat.com>
Date: Wed, 23 Dec 2020 15:54:01 -0800
Subject: [PATCH] add command to connect to the controller ssh via https...

This will be followed by an gentoken command or similar that uses
this as a jump host command..
---
 bitelab/__main__.py | 60 ++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 54 insertions(+), 6 deletions(-)

diff --git a/bitelab/__main__.py b/bitelab/__main__.py
index 1b05544..caa6b9f 100644
--- a/bitelab/__main__.py
+++ b/bitelab/__main__.py
@@ -126,16 +126,16 @@ async def fwd_data(reader, writer):
 
 		await writer.drain()
 
-async def run_exec(baseurl, authkey, board, args):
-	url = urllib.parse.urljoin(baseurl, 'board/%s/exec' %
-	    urllib.parse.quote(board, safe=''))
+async def run_exec(baseurl, path, args, authkey=None):
+	url = urllib.parse.urljoin(baseurl, path)
 	url = convert_to_ws(url)
 	stdin, stdout = await aioconsole.stream.get_standard_streams()
 
 	async with websockets.connect(url) as ws, wsfwd.WSFWDClient(ws.recv,
 	    ws.send) as client:
 		try:
-			await client.auth(dict(bearer=authkey))
+			if authkey is not None:
+				await client.auth(dict(bearer=authkey))
 
 			proc = await client.exec(args=args)
 
@@ -182,6 +182,9 @@ async def real_main():
 	parser_set.add_argument('board', type=str,
 	    help='name of the board or class')
 
+	parser_auth = subparsers.add_parser('contssh',
+	    help='open ssh session to the controller (internal)')
+
 	parser_exec = subparsers.add_parser('exec',
 	    help='run a program in the jail for a board')
 	parser_exec.add_argument('board', type=str,
@@ -197,8 +200,12 @@ async def real_main():
 	authkey = os.environ['BITELAB_AUTH']
 
 	if args.subparser_name == 'exec':
-		await run_exec(baseurl, authkey, args.board,
-		    [ args.prog ] + args.args)
+		path = 'board/%s/exec' % urllib.parse.quote(args.board, safe='')
+		await run_exec(baseurl, path=path, authkey=authkey,
+		    args=[ args.prog ] + args.args)
+		sys.exit(0) #pragma: no cover
+	elif args.subparser_name == 'contssh':
+		await run_exec(baseurl, path='ssh', args=[])
 		sys.exit(0) #pragma: no cover
 
 	client = AsyncClient(base_url=baseurl)
@@ -381,6 +388,47 @@ class TestExecClient(unittest.IsolatedAsyncioTestCase):
 
 		self.assertEqual(ret, 1)
 
+	@wsfwd.timeout(2)
+	async def test_contssh(self):
+		class TestServer(wsfwd.WSFWDCommon):
+			async def echo_handler(self, stream, msg):
+				self.sendstream(stream, msg)
+				await self.drain(stream)
+
+			async def handle_auth(self, msg):
+				assert msg['auth']['bearer'] == 'thisisanapikey'
+
+			async def handle_chanclose(self, msg):
+				self.add_tasks(asyncio.create_task(self.sendcmd(
+				    dict(cmd='chanclose',
+				    chan=self._stdout_stream))))
+
+			async def handle_exec(self, msg):
+				self._stdout_stream = msg['stdout']
+				self.add_stream_handler(msg['stdin'],
+				    functools.partial(self.echo_handler,
+				    msg['stdout']))
+
+				# pretend it's done immediately
+				self.add_tasks(asyncio.create_task(self.sendcmd(
+				    dict(cmd='exit', code=0))))
+
+		server = TestServer(self.toserver.get, self.toclient.put)
+
+		with patch.dict(sys.__dict__, dict(argv=[ 'rand', 'contssh', ])), \
+		    patch('websockets.connect') as webcon:
+			self.setup_websockets_mock(webcon)
+
+			inpdata = bytes(range(0, 255))
+
+			ret, stdout = await self.runAsyncMain(stdin=inpdata)
+
+		await server.__aexit__(None, None, None)
+
+		self.assertEqual(stdout, inpdata)
+
+		self.assertEqual(ret, 0)
+
 	@wsfwd.timeout(2)
 	async def test_exec(self):
 		class TestServer(wsfwd.WSFWDCommon):