Browse Source

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..
main
John-Mark Gurney 3 years ago
parent
commit
523132885f
1 changed files with 54 additions and 6 deletions
  1. +54
    -6
      bitelab/__main__.py

+ 54
- 6
bitelab/__main__.py View File

@@ -126,16 +126,16 @@ async def fwd_data(reader, writer):


await writer.drain() 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) url = convert_to_ws(url)
stdin, stdout = await aioconsole.stream.get_standard_streams() stdin, stdout = await aioconsole.stream.get_standard_streams()


async with websockets.connect(url) as ws, wsfwd.WSFWDClient(ws.recv, async with websockets.connect(url) as ws, wsfwd.WSFWDClient(ws.recv,
ws.send) as client: ws.send) as client:
try: try:
await client.auth(dict(bearer=authkey))
if authkey is not None:
await client.auth(dict(bearer=authkey))


proc = await client.exec(args=args) proc = await client.exec(args=args)


@@ -182,6 +182,9 @@ async def real_main():
parser_set.add_argument('board', type=str, parser_set.add_argument('board', type=str,
help='name of the board or class') 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', parser_exec = subparsers.add_parser('exec',
help='run a program in the jail for a board') help='run a program in the jail for a board')
parser_exec.add_argument('board', type=str, parser_exec.add_argument('board', type=str,
@@ -197,8 +200,12 @@ async def real_main():
authkey = os.environ['BITELAB_AUTH'] authkey = os.environ['BITELAB_AUTH']


if args.subparser_name == 'exec': 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 sys.exit(0) #pragma: no cover


client = AsyncClient(base_url=baseurl) client = AsyncClient(base_url=baseurl)
@@ -381,6 +388,47 @@ class TestExecClient(unittest.IsolatedAsyncioTestCase):


self.assertEqual(ret, 1) 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) @wsfwd.timeout(2)
async def test_exec(self): async def test_exec(self):
class TestServer(wsfwd.WSFWDCommon): class TestServer(wsfwd.WSFWDCommon):


Loading…
Cancel
Save