def init_key(self):
if self.ssh_custom_key:
return
for filename in os.listdir(self.config.get('local', 'data_dir')):
if not filename.startswith('yakey') or not os.path.isfile(
os.path.join(self.config.get('local', 'data_dir'),
filename)):
continue
key_path = os.path.join(self.config.get('local', 'data_dir'),
filename)
self.key_name = key_path.split(os.sep)[-1]
pmk_key = RSAKey.from_private_key_file(key_path)
logging.info('LOADED KEY %s' % key_path)
break
else:
self.key_name = self.get_rnd_name('yakey')
key = rsa.generate_private_key(backend=crypto_default_backend(),
public_exponent=65537,
key_size=2048)
pmk_key = RSAKey(key=key)
key_path = os.path.join(self.config.get('local', 'data_dir'),
self.key_name)
pmk_key.write_private_key_file(key_path)
logging.info('WRITTEN KEY %s' % key_path)
self.public_key = "%s %s" % (pmk_key.get_name(), pmk_key.get_base64())
self.ssh_custom_key = {'pkey': pmk_key}
if self.yascheduler:
self.yascheduler.ssh_custom_key = self.ssh_custom_key
class KeyPair():
KEY_SIZE = 4096
SSH_PUB_KEY_PREFIX = 'ssh-rsa '
def __init__(self, private_key=None, private_key_path=None, pub_data=None, _key=None, password=None):
if private_key:
self._key = RSAKey(file_obj=io.StringIO(private_key))
elif private_key_path:
self._key = RSAKey(filename=private_key_path, password=password)
elif pub_data:
if pub_data.startswith(self.SSH_PUB_KEY_PREFIX):
pub_data = pub_data[len(self.SSH_PUB_KEY_PREFIX):]
self._key = RSAKey(data=b64decode(pub_data.encode('utf-8')))
elif _key:
self._key = _key
else:
self._key = RSAKey.generate(self.KEY_SIZE)
def public_key_str(self):
return self.SSH_PUB_KEY_PREFIX + self._key.get_base64()
def private_key_str(self):
buf = io.StringIO()
self._key.write_private_key(buf)
buf.seek(0)
private = buf.read()
return private
def get_keypair_str(self):
return self.public_key_str(), self.private_key_str()
def as_paramiko(self):
return self._key
def as_pycrypto(self):
if self._key.can_sign():
return RSA.importKey(self.private_key_str())
else:
return RSA.importKey(self.public_key_str())
def encrypt(self, data, encoding='utf-8', encode_payload=False):
if isinstance(data, str) and encoding:
data = data.encode(encoding)
cipher = PKCS1_OAEP.new(self.as_pycrypto())
ciphertext = cipher.encrypt(data)
if encode_payload:
return base64.b64encode(ciphertext).decode('utf-8')
else:
return ciphertext
def decrypt(self, ciphertext, encoding='utf-8'):
if isinstance(ciphertext, str):
ciphertext = base64.b64decode(ciphertext.encode('utf-8'))
cipher = PKCS1_OAEP.new(self.as_pycrypto())
data = cipher.decrypt(ciphertext)
if encoding:
text = data.decode(encoding)
return text
def to_file(self, path, password=None):
if self._key.can_sign():
self._key.write_private_key_file(path, password=password)
