testing/legion/ssl_util.py

   1 # Copyright 2015 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Utility library to add SSL support to the RPC server."""
   6
   7 import logging
   8 import ssl
   9 import subprocess
  10 import tempfile
  11
  12 #pylint: disable=relative-import
  13 import common_lib
  14 import jsonrpclib
  15 import SimpleJSONRPCServer
  16
  17
  18 class Error(Exception):
  19   pass
  20
  21
  22 def CreateKeyFile():
  23   """Creates an SSL keyfile and returns the path."""
  24   keyfile = tempfile.mkstemp()[1]
  25   cmd = [
  26     'openssl',
  27     'genrsa',
  28     '-out', keyfile,
  29     '2048'
  30     ]
  31   _RunCommand(cmd)
  32   return keyfile
  33
  34
  35 def CreateCsrFile(keyfile):
  36   """Creates an SSL CSR file and returns the path."""
  37   csrfile = tempfile.mkstemp()[1]
  38   cmd = [
  39       'openssl',
  40       'req',
  41       '-new',
  42       '-key', keyfile,
  43       '-out', csrfile,
  44       '-subj', '/C=NA/ST=NA/L=NA/O=Chromium/OU=Test/CN=chromium.org'
  45       ]
  46   _RunCommand(cmd)
  47   return csrfile
  48
  49
  50 def CreateCrtFile(keyfile, csrfile):
  51   """Creates an SSL CRT file and returns the path."""
  52   crtfile = tempfile.mkstemp()[1]
  53   cmd = [
  54       'openssl',
  55       'x509',
  56       '-req',
  57       '-days', '1',
  58       '-in', csrfile,
  59       '-signkey', keyfile,
  60       '-out', crtfile
  61       ]
  62   _RunCommand(cmd)
  63   return crtfile
  64
  65
  66 def CreatePemFile():
  67   """Creates an SSL PEM file and returns the path."""
  68   keyfile = CreateKeyFile()
  69   csrfile = CreateCsrFile(keyfile)
  70   crtfile = CreateCrtFile(keyfile, csrfile)
  71   pemfile = tempfile.mkstemp()[1]
  72   with open(keyfile) as k:
  73     with open(crtfile) as c:
  74       with open(pemfile, 'wb') as p:
  75         p.write('%s\n%s' % (k.read(), c.read()))
  76   return pemfile
  77
  78
  79 def _RunCommand(cmd):
  80   try:
  81     p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  82   except OSError as e:
  83     raise Error('Failed to run %s: %s' % (' '.join(cmd), e))
  84   out, err = p.communicate()
  85   if p.returncode != 0:
  86     raise Error(err)
  87   return out
  88
  89
  90 class SslRpcServer(SimpleJSONRPCServer.SimpleJSONRPCServer):
  91   """Class to add SSL support to the RPC server."""
  92
  93   def __init__(self, *args, **kwargs):
  94     SimpleJSONRPCServer.SimpleJSONRPCServer.__init__(self, *args, **kwargs)
  95     self.socket = ssl.wrap_socket(self.socket, certfile=CreatePemFile(),
  96                                   server_side=True)
  97
  98   @staticmethod
  99   def Connect(server, port=common_lib.SERVER_PORT):
 100     """Creates and returns a connection to an SSL RPC server."""
 101     addr = 'https://%s:%d' % (server, port)
 102     logging.debug('Connecting to RPC server at %s', addr)
 103     return jsonrpclib.ServerProxy(addr, allow_none=True)