common/python/rift/mano/utils/ssh_keys.py

   1 #
   2 # Licensed under the Apache License, Version 2.0 (the "License"); you may
   3 # not use this file except in compliance with the License. You may obtain
   4 # a copy of the License at
   5 #
   6 # http://www.apache.org/licenses/LICENSE-2.0
   7 #
   8 # Unless required by applicable law or agreed to in writing, software
   9 # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  10 # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  11 # License for the specific language governing permissions and limitations
  12 # under the License.
  13
  14 # Copyright 2016 RIFT.io Inc
  15
  16
  17 import argparse
  18 import logging
  19 import os
  20 import socket
  21 import stat
  22 import sys
  23 import tempfile
  24
  25 from Crypto.PublicKey import RSA
  26
  27
  28 class ManoSshKey(object):
  29     '''
  30     Generate a SSH key pair and store them in a file
  31     '''
  32
  33     def __init__(self, log, size=2048):
  34         self._log = log
  35         self._size = size
  36
  37         self._key = None
  38         self._key_pem = None
  39         self._pub_ssh = None
  40         self._key_file = None
  41         self._pub_file = None
  42
  43     @property
  44     def log(self):
  45         return self._log
  46
  47     @property
  48     def size(self):
  49         return self._size
  50
  51     @property
  52     def private_key(self):
  53         if self._key is None:
  54             self._gen_keys()
  55         return self._key_pem
  56
  57     @property
  58     def public_key(self):
  59         if self._key is None:
  60             self._gen_keys()
  61         return self._pub_ssh
  62
  63     @property
  64     def private_key_file(self):
  65         return self._key_file
  66
  67     @property
  68     def public_key_file(self):
  69         return self._pub_file
  70
  71     def _gen_keys(self):
  72         if self._key:
  73             return
  74
  75         self.log.info("Generating key of size: {}".format(self.size))
  76
  77         self._key = RSA.generate(self.size, os.urandom)
  78         self._key_pem = self._key.exportKey('PEM').decode('utf-8')
  79         self.log.debug("Private key PEM: {}".format(self._key_pem))
  80
  81         # Public key export as 'OpenSSH' has a bug
  82         # (https://github.com/dlitz/pycrypto/issues/99)
  83
  84         username = None
  85         try:
  86             username = os.getlogin()
  87             hostname = socket.getfqdn()
  88         except OSError:
  89             pass
  90
  91         pub = self._key.publickey().exportKey('OpenSSH').decode('utf-8')
  92         if username:
  93             self._pub_ssh = '{} {}@{}'.format(pub, username, hostname)
  94         else:
  95             self._pub_ssh = pub
  96         self.log.debug("Public key SSH: {}".format(self._pub_ssh))
  97
  98     def write_to_disk(self,
  99                       name="id_rsa",
 100                       directory="."):
 101         if self._key is None:
 102             self._gen_keys()
 103
 104         path = os.path.abspath(directory)
 105         self._pub_file = "{}/{}.pub".format(path, name)
 106         self._key_file = "{}/{}.key".format(path, name)
 107
 108         with open(self._key_file, 'w') as content_file:
 109             content_file.write(self.private_key)
 110             os.chmod(self._key_file, stat.S_IREAD|stat.S_IWRITE)
 111
 112         with open(self._pub_file, 'w') as content_file:
 113             content_file.write(self.public_key)
 114
 115 if __name__ == "__main__":
 116     parser = argparse.ArgumentParser(description='Generate SSH key pair')
 117     parser.add_argument("-s", "--size", type=int, default=2048, help="Key size")
 118     parser.add_argument("-d", "--directory", help="Directory to store the keys")
 119     parser.add_argument("-n", "--name", help="Name for the key file")
 120     parser.add_argument("--debug", help="Enable debug logging",
 121                         action="store_true")
 122     args = parser.parse_args()
 123
 124     fmt = logging.Formatter(
 125         '%(asctime)-23s %(levelname)-5s  (%(name)s@%(process)d:' \
 126         '%(filename)s:%(lineno)d) - %(message)s')
 127     stderr_handler = logging.StreamHandler(stream=sys.stderr)
 128     stderr_handler.setFormatter(fmt)
 129     if args.debug:
 130         logging.basicConfig(level=logging.DEBUG)
 131     else:
 132         logging.basicConfig(level=logging.INFO)
 133     log = logging.getLogger('rw-mano-ssh-keys')
 134     log.addHandler(stderr_handler)
 135
 136     log.info("Args passed: {}".format(args))
 137     if args.directory:
 138         path = args.directory
 139     else:
 140         path = tempfile.mkdtemp()
 141
 142     kp = ManoSshKey(log, size=args.size)
 143     kp.write_to_disk(directory=path)
 144     log.info("Private Key: {}".format(kp.private_key))
 145     log.info("Public key: {}".format(kp.public_key))
 146     log.info("Key file: {}, Public file: {}".format(kp.private_key_file,
 147                                                     kp.public_key_file))