bitelab/bitelab/__main__.py

#
# Copyright (c) 2020 The FreeBSD Foundation
#
# This software1 was developed by John-Mark Gurney under sponsorship
# from the FreeBSD Foundation.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1.  Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
# 2.  Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

from httpx import AsyncClient, Auth
from io import StringIO
from starlette.status import HTTP_200_OK
from starlette.status import HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED, \
    HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND, HTTP_409_CONFLICT
from unittest.mock import patch, AsyncMock, Mock, mock_open

from . import BiteAuth, Board

import argparse
import asyncio
import contextlib
import io
import os
import sys
import unittest
import urllib

def check_res_code(res):
	if res.status_code == HTTP_401_UNAUTHORIZED:
		print('Invalid authentication credentials.')
		sys.exit(1)
	elif res.status_code != HTTP_200_OK:
		print('Got status: %d, json: %s' % (res.status_code, res.json()))
		sys.exit(1)

def makebool(s):
	s = s.lower()
	if s in { 'on', 'true', '1' }:
		return True
	elif s in { 'off', 'false', '0' }:
		return False

	raise ValueError('invalid value for bool: %s' % repr(s))

_typeconv = dict(power=makebool)
def make_attrs(*args):
	'''Convert list of attributes to a dictionary.'''

	r = {}
	for i in args:
		k, v = i.split('=', 1)
		r[k] = _typeconv[k](v)

	return r

_httpxargs = dict(timeout=20)

def output_board(brd):
	print('Name:\t%s' % brd.name)
	print('Class:\t%s' % brd.brdclass)
	print('Attributes:')
	for i in sorted(brd.attrs):
		print('\t%s\t%s' % (i, repr(brd.attrs[i])))

def get_sshpubkey(fname):
	if fname is not None:
		with open(fname) as fp:
			return fp.read()

	for i in ('id_ed25519', 'id_rsa'):
		fname = os.path.expanduser(os.path.join('~', '.ssh', i + '.pub'))
		with contextlib.suppress(IOError), open(fname) as fp:
			return fp.read()

	raise IOError

async def real_main():
	baseurl = os.environ['BITELAB_URL']
	authkey = os.environ['BITELAB_AUTH']

	client = AsyncClient(base_url=baseurl)

	parser = argparse.ArgumentParser()
	subparsers = parser.add_subparsers(title='subcommands',
	    dest='subparser_name',
	    description='valid subcommands', help='additional help')

	parse_list = subparsers.add_parser('list', help='list available board classes')

	parser_reserve = subparsers.add_parser('reserve', aliases=[ 'release' ], help='reserve/release a board')
	parser_reserve.add_argument('-i', metavar='identity_file', type=str, help='file name for ssh public key')
	parser_reserve.add_argument('board', type=str, help='name of the board or class')

	parser_set = subparsers.add_parser('set', help='set attributes on a board')
	parser_set.add_argument('setvars', type=str, nargs='+', help='name of the board or class')
	parser_set.add_argument('board', type=str, help='name of the board or class')

	args = parser.parse_args()
	#print(repr(args), file=sys.stderr)

	try:
		if args.subparser_name == 'list':
			res = await client.get('board/classes',
			    auth=BiteAuth(authkey), **_httpxargs)

			check_res_code(res)

			print('Classes:')
			for i in res.json():
				print('\t' + i)

			res.close()
		elif args.subparser_name in ('reserve', 'release'):
			kwargs = _httpxargs.copy()
			with contextlib.suppress(IOError):
				kwargs['json'] = dict(sshpubkey=get_sshpubkey(args.i))
			res = await client.post('board/%s/%s' %
			    (urllib.parse.quote(args.board, safe=''),
			     args.subparser_name),
			    auth=BiteAuth(authkey), **kwargs)

			check_res_code(res)

			brd = Board.parse_obj(res.json())
			output_board(brd)

		elif args.subparser_name == 'set':
			board_id = sys.argv[-1]
			res = await client.post('board/%s/attrs' % 
			    urllib.parse.quote(board_id, safe=''),
			    auth=BiteAuth(authkey),
			    json=make_attrs(*sys.argv[2:-1]), **_httpxargs)

			check_res_code(res)

			brd = Board.parse_obj(res.json())
			output_board(brd)

	finally:
		await client.aclose()

def main():
	asyncio.run(real_main())

if __name__ == '__main__': #pragma: no cover
	main()

@patch.dict(os.environ, dict(BITELAB_URL='http://someserver/'))
@patch.dict(os.environ, dict(BITELAB_AUTH='thisisanapikey'))
class TestClient(unittest.TestCase):
	def setUp(self):
		self.ac_patcher = patch(__name__ + '.AsyncClient')
		self.ac = self.ac_patcher.start()
		self.addCleanup(self.ac_patcher.stop)

		self.acg = self.ac.return_value.get = AsyncMock()
		self.acaclose = self.ac.return_value.aclose = AsyncMock()
		self.acgr = self.acg.return_value = Mock()

		self.acp = self.ac.return_value.post = AsyncMock()
		self.acpr = self.acp.return_value = Mock()

	def runMain(self):
		try:
			stdout = StringIO()
			with patch.dict(sys.__dict__, dict(stdout=stdout)):
				main()

			ret = 0
		except SystemExit as e:
			ret = e.code

		return ret, stdout.getvalue()

	def test_sshpubkey(self):
		fname = 'fname'
		rdata = 'foo'
		with patch('builtins.open', mock_open(read_data=rdata)) as mock_file:
			self.assertEqual(get_sshpubkey(fname), rdata)
			mock_file.assert_called_with(fname)

		with patch('builtins.open') as mock_file:
			mock_file.side_effect = [ IOError(), mock_open(read_data=rdata)() ]
			self.assertEqual(get_sshpubkey(None), rdata)

			for i in ('id_ed25519', 'id_rsa'):
				fname = os.path.expanduser(os.path.join('~', '.ssh', i + '.pub'))
				mock_file.assert_any_call(fname)

	@patch.dict(sys.__dict__, dict(argv=[ '', 'list' ]))
	def test_list_failure(self):
		ac = self.ac
		acg = self.acg
		acg.return_value.status_code = HTTP_401_UNAUTHORIZED
		acg.return_value.json.return_value = {
		    'detail': 'Invalid authentication credentials'
		}

		ret, stdout = self.runMain()

		output = '''Invalid authentication credentials.
'''

		self.assertEqual(ret, 1)
		# XXX -- really should go to stderr
		self.assertEqual(stdout, output)

		ac.assert_called_with(base_url='http://someserver/')

		acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'), **_httpxargs)

	# XXX - add error cases for UI

	@patch.dict(sys.__dict__, dict(argv=[ '', 'list' ]))
	def test_list(self):
		ac = self.ac
		acg = self.acg
		acg.return_value.status_code = HTTP_200_OK
		acg.return_value.json.return_value = { 'cora-z7s': {
		    'arch': 'arm-armv7', 'clsname': 'cora-z7s', }}

		ret, stdout = self.runMain()

		output = '''Classes:
	cora-z7s
'''

		self.assertEqual(ret, 0)
		self.assertEqual(stdout, output)

		ac.assert_called_with(base_url='http://someserver/')

		acg.assert_called_with('board/classes', auth=BiteAuth('thisisanapikey'), **_httpxargs)

	# XXX - add error cases for UI

	@patch('bitelab.__main__.get_sshpubkey')
	@patch.dict(sys.__dict__, dict(argv=[ '', 'reserve', 'cora-z7s' ]))
	def test_reserve(self, gspk):
		gspk.side_effect = IOError()
		ac = self.ac
		acp = self.acp
		acp.return_value.status_code = HTTP_200_OK
		acp.return_value.json.return_value = Board(name='cora-1',
		    brdclass='cora-z7s', reserved=True,
		    attrs={
		        'ip': '172.20.20.5',
		        'power': False,
		    }).dict()

		keydata = 'randomkeydata'

		ret, stdout = self.runMain()

		output = '''Name:\tcora-1
Class:\tcora-z7s
Attributes:
\tip\t'172.20.20.5'
\tpower\tFalse
'''

		self.assertEqual(ret, 0)
		self.assertEqual(stdout, output)

		ac.assert_called_with(base_url='http://someserver/')

		#args = { 'sshpubkey': keydata }
		acp.assert_called_with('board/cora-z7s/reserve',
		    #json=args,
		    auth=BiteAuth('thisisanapikey'), **_httpxargs)

	@patch('bitelab.__main__.get_sshpubkey')
	@patch.dict(sys.__dict__, dict(argv=[ '', 'reserve', '-i', 'fixtures/testsshkey.pub', 'cora-z7s' ]))
	def test_reserve_ssh(self, gspk):
		ac = self.ac
		acp = self.acp
		acp.return_value.status_code = HTTP_200_OK
		acp.return_value.json.return_value = Board(name='cora-1',
		    brdclass='cora-z7s', reserved=True,
		    attrs={
		        'ip': '172.20.20.5',
		        'power': False,
		    }).dict()

		keydata = 'keydata'
		gspk.return_value = keydata

		ret, stdout = self.runMain()

		output = '''Name:\tcora-1
Class:\tcora-z7s
Attributes:
\tip\t'172.20.20.5'
\tpower\tFalse
'''

		self.assertEqual(ret, 0)
		self.assertEqual(stdout, output)

		ac.assert_called_with(base_url='http://someserver/')

		acp.assert_called_with('board/cora-z7s/reserve',
		    json=dict(sshpubkey=keydata),
		    auth=BiteAuth('thisisanapikey'), **_httpxargs)

	@patch('bitelab.__main__.get_sshpubkey')
	@patch.dict(sys.__dict__, dict(argv=[ '', 'release', 'cora-z7s' ]))
	def test_release(self, gspk):
		gspk.side_effect = IOError
		ac = self.ac
		acp = self.acp
		acp.return_value.status_code = HTTP_200_OK
		acp.return_value.json.return_value = Board(name='cora-1',
		    brdclass='cora-z7s', reserved=False,
		    attrs={
		        'power': False,
		    }).dict()

		ret, stdout = self.runMain()

		output = '''Name:\tcora-1
Class:\tcora-z7s
Attributes:
\tpower\tFalse
'''

		self.assertEqual(ret, 0)
		self.assertEqual(stdout, output)

		ac.assert_called_with(base_url='http://someserver/')

		acp.assert_called_with('board/cora-z7s/release',
		    auth=BiteAuth('thisisanapikey'), **_httpxargs)

	@patch('bitelab.__main__._typeconv', dict(power=makebool, other=makebool))
	@patch.dict(sys.__dict__, dict(argv=[ '', 'set', 'power=on', 'other=off', 'cora-z7s' ]))
	def test_set(self):
		ac = self.ac
		acp = self.acp
		acp.return_value.status_code = HTTP_200_OK
		acp.return_value.json.return_value = Board(name='cora-1',
		    brdclass='cora-z7s', reserved=True,
		    attrs={
		        'ip': '172.20.20.5',
		        'power': True,
		    }).dict()

		ret, stdout = self.runMain()

		output = '''Name:\tcora-1
Class:\tcora-z7s
Attributes:
\tip\t'172.20.20.5'
\tpower\tTrue
'''

		self.assertEqual(ret, 0)
		self.assertEqual(stdout, output)

		ac.assert_called_with(base_url='http://someserver/')

		acp.assert_called_with('board/cora-z7s/attrs',
		    auth=BiteAuth('thisisanapikey'), json=dict(power=True, other=False),
		    **_httpxargs)

	def test_make_attrs(self):
		self.assertEqual(make_attrs('power=on'), dict(power=True))
		self.assertEqual(make_attrs('power=True'), dict(power=True))
		self.assertEqual(make_attrs('power=1'), dict(power=True))
		self.assertEqual(make_attrs('power=off'), dict(power=False))
		self.assertEqual(make_attrs('power=FALSE'), dict(power=False))
		self.assertEqual(make_attrs('power=0'), dict(power=False))

		self.assertRaises(ValueError, make_attrs, 'power=bar')

	def test_check_res_code(self):
		# XXX - on weird errors, res.json() will fail w/ json.decoder.JSONDecodeError
		pass