|
|
@@ -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): |
|
|
|