def setUpClass(cls):
# Prepare server
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(('127.0.0.1', 12345))
sock.listen(10)
cls.server_sock = sock
th = threading.Thread(target=cls.__server)
th.setDaemon(True)
th.start()
cls.server_thread = th
cls.server_ready = threading.Event()
cls.server_end = threading.Event()
# Prepare BackendConnector
cfg = GlobalConfigStore()
cfg.agent_username = "******"
cfg.agentkey = RSAKey.generate(1024)
cfg.backend_hostkey = RSAKey.generate(1024)
cfg.host = "127.0.0.1"
cfg.port = 12345
cls.ingress = queue.Queue()
cls.egress = queue.Queue()
cls.connector = BackendConnector()
endpoint = BackendConnector.EndPoint(ingress=cls.ingress,
egress=cls.egress)
cls.connector.register(endpoint)
cls.connector.start()
def bastio_main():
"""Main application entry point."""
signal.signal(signal.SIGINT, __sig_handler)
signal.signal(signal.SIGTERM, __sig_handler)
cfg = GlobalConfigStore()
# Parse command line arguments
cmd = CommandLine()
command = cmd.parse()
if command == 'generate-key':
try:
key = RSAKey.generate(cfg.bits)
key.write_private_key_file(cfg.agentkey)
except Exception as ex:
_die(ex.message)
_die("generated {}-bit key successfully".format(cfg.bits), True)
elif command == 'upload-key':
try:
agentkey = RSAKey.from_private_key_file(cfg.agentkey)
except Exception as ex:
_die(ex.message)
try:
# Replace old key
if cfg.new_agentkey:
new_agentkey = RSAKey.from_private_key_file(cfg.new_agentkey)
upload_public_key(cfg.apikey, new_agentkey.get_public_key(),
agentkey.get_public_key())
else:
upload_public_key(cfg.apikey, agentkey.get_public_key())
except Exception as ex:
_die(ex.message)
_die("uploaded public key successfully", True)
elif command == 'start':
try:
cfg.agent_username = cfg.apikey
cfg.agentkey = RSAKey.from_private_key_file(cfg.agentkey)
cfg.backend_hostkey = download_backend_hostkey()
except Exception as ex:
_die(ex.message)
### All that is below is part of the ``start`` command
# Set logging based on debug status
if cfg.debug:
Logger().enable_stream()
else:
Logger().enable_syslog()
cfg.threadpool = GlobalThreadPool(cfg.minthreads)
cfg.processor = Processor()
cfg.connector = BackendConnector()
cfg.connector.register(cfg.processor.endpoint())
cfg.connector.start()
signal.pause()
def upload_public_key(api_key, public_key, old_public_key=None):
"""Upload agent's public key to Bastio on the account specified by ``api_key``.
This action will create a new server when ``old_public_key`` is not specified
and it is the first time we see this ``public_key``. However the public key
we store that identifies this server will be replaced IFF we already have
``old_public_key`` in our records and ``public_key`` does not exist in our
database.
:param api_key:
The API key for the Bastio account.
:type api_key:
str
:param public_key:
The agent's public key to be uploaded to Bastio's servers.
:type public_key:
The string output of :func:`bastio.ssh.crypto.RSAKey.get_public_key`.
:param old_public_key:
The agent's old public key to be replaced by a new one.
:type old_public_key:
The string output of :func:`bastio.ssh.crypto.RSAKey.get_public_key`.
:raises:
:class:`bastio.excepts.BastioAccountError`
"""
errmsg = "upload public key failed: "
if not old_public_key:
old_public_key = ''
if old_public_key and not RSAKey.validate_public_key(old_public_key):
raise BastioAccountError(errmsg + "invalid old public key")
if not RSAKey.validate_public_key(public_key):
raise BastioAccountError(errmsg + "invalid new public key")
payload = Json()
payload.api_key = api_key
payload.public_key = public_key
payload.old_public_key = old_public_key
headers = {'Content-type': 'application/json'}
response = __send_request('post', url=__upload_key_endpoint, verify=True,
data=payload.to_json(), headers=headers)
if response.status_code == requests.codes.bad: # 400
raise BastioAccountError(errmsg + "missing or invalid field")
elif response.status_code == requests.codes.forbidden: # 403
raise BastioAccountError(errmsg + "not authorized or invalid API key")
elif response.status_code == requests.codes.okay: # 200
return
# An unexpected response status code
raise BastioAccountError(
errmsg + "unexpected response status code ({})".format(
response.status_code))
def download_backend_hostkey():
"""Get Bastio's backend SSH host key.
:returns:
:class:`bastio.ssh.crypto.RSAKey`
:raises:
:class:`bastio.excepts.BastioAccountError`
"""
errmsg = "get backend host key failed: "
response = __send_request('get', url=__download_hostkey_endpoint, verify=True)
if response.status_code != requests.codes.okay: # 200
raise BastioAccountError(errmsg + "unable to retrieve backend's host key")
public_key = response.json()['payload']
if not RSAKey.validate_public_key(public_key):
raise BastioAccountError(errmsg + "invalid host key")
return RSAKey.from_public_key(public_key)
def parse(cls, obj, traverse=True):
"""Check remove-key message fields and validate them.
Return a new object of type ``cls`` containing the validated
action object.
:param obj:
A JSON object containing the relevant fields for this action message.
:type obj:
:class:`bastio.mixin.Json`
:param traverse:
Whether to traverse ``parse`` on all the classes in the hierarchy.
:type traverse:
bool
:returns:
A new object of type ``cls`` containing the validated action
object.
"""
if 'public_key' not in obj:
raise BastioMessageError("public_key field is missing")
if not RSAKey.validate_public_key(obj.public_key):
raise BastioMessageError("public_key field is invalid")
if traverse:
return super(RemoveKeyMessage, cls).parse(obj)
def test_key_loading(self):
pubkey = self.key.get_public_key()
self.assertIsInstance(RSAKey.from_public_key(pubkey), RSAKey)
def test_key_validation(self):
priv = self.key.get_private_key()
self.assertTrue(RSAKey.validate_private_key(priv))
self.assertTrue(RSAKey.validate_private_key_file('TEST_PRIVATE_KEY'))
self.assertTrue(RSAKey.validate_public_key(self.key.get_public_key()))
def test_key_generation(self):
# Test invalid key size
with self.assertRaises(BastioCryptoError):
RSAKey.generate(1111)
self.assertTrue(self.key.get_private_key())
self.assertTrue(self.key.get_public_key())
def setUpClass(cls):
cls.key = RSAKey.generate(1024)
cls.key_file = 'TEST_PRIVATE_KEY'
file(cls.key_file, 'wb').write(cls.key.get_private_key())
def setUpClass(cls):
cls._public_key = RSAKey.generate(1024).get_public_key()
cls._proc = Processor()
def parse(self):
"""Parse and validate arguments from the command line and set
global configurations.
"""
self.args = self.parser.parse_args()
cfg = GlobalConfigStore()
cfg.prog = self.parser.prog
cfg.debug = self.args.debug
# Check of configuration file is available to us
conf_avail = False
if self.args.config:
try:
cfg.load(self.args.config)
conf_avail = True
except BastioConfigError as ex:
self.parser.error(ex.message)
# Check and validate agent's key if we are about to upload the key to
# Bastio's servers or we are about to start the agent
if self.args.command in ('upload-key', 'start'):
# Get agent key file path from configuration file (if available)
# or from the command line argument
try:
if conf_avail:
cfg.apikey = cfg.apikey if cfg.get_apikey else \
self.args.api_key
cfg.agentkey = cfg.agentkey if cfg.get_agentkey else \
self.args.agent_key
else:
cfg.apikey = self.args.api_key
cfg.agentkey = self.args.agent_key
except BastioConfigError as ex:
_die(ex.message)
# Check agent's key file readability and validate it
res = _check_file_readability(cfg.agentkey)
if not res[0]:
self.parser.error('agent key file `{}` does not exist'.format(
cfg.agentkey))
if not res[1]:
self.parser.error(('permission to read the agent key file `{}` '
'is denied').format(cfg.agentkey))
res = RSAKey.validate_private_key_file(cfg.agentkey)
if not res:
self.parser.error('agent key file `{}` is invalid'.format(
cfg.agentkey))
# Parse and validate commands and their arguments
if self.args.command == 'generate-key':
try:
if conf_avail:
cfg.agentkey = cfg.agentkey if cfg.get_agentkey else \
self.args.agent_key
else:
cfg.agentkey = self.args.agent_key
cfg.bits = self.args.bits
except BastioConfigError as ex:
_die(ex.message)
elif self.args.command == 'upload-key':
try:
# Check new key file's readability and validate it if provided
new_key = self.args.new_agent_key
if new_key:
res = _check_file_readability(new_key)
if not res[0]:
self.parser.error(
'new agent key file `{}` does not exist'.format(
new_key))
if not res[1]:
self.parser.error((
'permission to read the new agent key file `{}` '
'is denied').format(new_key))
res = RSAKey.validate_private_key_file(new_key)
if not res:
self.parser.error(
'new agent key file `{}` is invalid'.format(
new_key))
cfg.new_agentkey = new_key
except BastioConfigError as ex:
_die(ex.message)
elif self.args.command == 'start':
try:
if conf_avail:
cfg.host = cfg.host if cfg.get_host else self.args.host
cfg.port = cfg.port if cfg.getint_port else self.args.port
cfg.stacksize = cfg.stacksize if cfg.getint_stacksize else \
self.args.stack_size
cfg.minthreads = cfg.minthreads if cfg.getint_minthreads else \
self.args.min_threads
else:
cfg.host = self.args.host
cfg.port = self.args.port
cfg.stacksize = self.args.stack_size
cfg.minthreads = self.args.min_threads
except BastioConfigError as ex:
_die(ex.message)
else:
# NOTE: This execution branch is blocked by argparse
# so it is here only to account for extremely unlikely cases
_die("unsupported command `{}`".format(self.args.command))
return self.args.command
def setUpClass(cls):
cls.pubkey = RSAKey.generate(1024).get_public_key()