michael@0: #!/usr/bin/env python
michael@0: # This Source Code Form is subject to the terms of the Mozilla Public
michael@0: # License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0: # file, You can obtain one at http://mozilla.org/MPL/2.0/.
michael@0: 
michael@0: # This script exists to generate the Certificate Authority and server
michael@0: # certificates used for SSL testing in Mochitest. The already generated
michael@0: # certs are located at $topsrcdir/build/pgo/certs/ .
michael@0: 
michael@0: import mozinfo
michael@0: import os
michael@0: import random
michael@0: import re
michael@0: import shutil
michael@0: import subprocess
michael@0: import sys
michael@0: import tempfile
michael@0: 
michael@0: from mozbuild.base import MozbuildObject
michael@0: from mozfile import NamedTemporaryFile
michael@0: from mozprofile.permissions import ServerLocations
michael@0: 
michael@0: dbFiles = [
michael@0:   re.compile("^cert[0-9]+\.db$"),
michael@0:   re.compile("^key[0-9]+\.db$"),
michael@0:   re.compile("^secmod\.db$")
michael@0: ]
michael@0: 
michael@0: def unlinkDbFiles(path):
michael@0:   for root, dirs, files in os.walk(path):
michael@0:     for name in files:
michael@0:       for dbFile in dbFiles:
michael@0:         if dbFile.match(name) and os.path.exists(os.path.join(root, name)):
michael@0:           os.unlink(os.path.join(root, name))
michael@0: 
michael@0: def dbFilesExist(path):
michael@0:   for root, dirs, files in os.walk(path):
michael@0:     for name in files:
michael@0:       for dbFile in dbFiles:
michael@0:         if dbFile.match(name) and os.path.exists(os.path.join(root, name)):
michael@0:           return True
michael@0:   return False
michael@0: 
michael@0: 
michael@0: def runUtil(util, args, inputdata = None):
michael@0:   env = os.environ.copy()
michael@0:   if mozinfo.os == "linux":
michael@0:     pathvar = "LD_LIBRARY_PATH"
michael@0:     app_path = os.path.dirname(util)
michael@0:     if pathvar in env:
michael@0:       env[pathvar] = "%s%s%s" % (app_path, os.pathsep, env[pathvar])
michael@0:     else:
michael@0:       env[pathvar] = app_path
michael@0:   proc = subprocess.Popen([util] + args, env=env,
michael@0:                           stdin=subprocess.PIPE if inputdata else None)
michael@0:   proc.communicate(inputdata)
michael@0:   return proc.returncode
michael@0: 
michael@0: 
michael@0: def createRandomFile(randomFile):
michael@0:   for count in xrange(0, 2048):
michael@0:     randomFile.write(chr(random.randint(0, 255)))
michael@0: 
michael@0: 
michael@0: def createCertificateAuthority(build, srcDir):
michael@0:   certutil = build.get_binary_path(what="certutil")
michael@0:   pk12util = build.get_binary_path(what="pk12util")
michael@0: 
michael@0:   #TODO: mozfile.TemporaryDirectory
michael@0:   tempDbDir = tempfile.mkdtemp()
michael@0:   with NamedTemporaryFile() as pwfile, NamedTemporaryFile() as rndfile:
michael@0:     pgoCAModulePathSrc = os.path.join(srcDir, "pgoca.p12")
michael@0:     pgoCAPathSrc = os.path.join(srcDir, "pgoca.ca")
michael@0: 
michael@0:     pwfile.write("\n")
michael@0: 
michael@0:     # Create temporary certification database for CA generation
michael@0:     status = runUtil(certutil, ["-N", "-d", tempDbDir, "-f", pwfile.name])
michael@0:     if status:
michael@0:       return status
michael@0: 
michael@0:     createRandomFile(rndfile)
michael@0:     status = runUtil(certutil, ["-S", "-d", tempDbDir, "-s", "CN=Temporary Certificate Authority, O=Mozilla Testing, OU=Profile Guided Optimization", "-t", "C,,", "-x", "-m", "1", "-v", "120", "-n", "pgo temporary ca", "-2", "-f", pwfile.name, "-z", rndfile.name], "Y\n0\nN\n")
michael@0:     if status:
michael@0:       return status
michael@0: 
michael@0:     status = runUtil(certutil, ["-L", "-d", tempDbDir, "-n", "pgo temporary ca", "-a", "-o", pgoCAPathSrc, "-f", pwfile.name])
michael@0:     if status:
michael@0:       return status
michael@0: 
michael@0:     status = runUtil(pk12util, ["-o", pgoCAModulePathSrc, "-n", "pgo temporary ca", "-d", tempDbDir, "-w", pwfile.name, "-k", pwfile.name])
michael@0:     if status:
michael@0:       return status
michael@0: 
michael@0:   shutil.rmtree(tempDbDir)
michael@0:   return 0
michael@0: 
michael@0: 
michael@0: def createSSLServerCertificate(build, srcDir):
michael@0:   certutil = build.get_binary_path(what="certutil")
michael@0:   pk12util = build.get_binary_path(what="pk12util")
michael@0: 
michael@0:   with NamedTemporaryFile() as pwfile, NamedTemporaryFile() as rndfile:
michael@0:     pgoCAPath = os.path.join(srcDir, "pgoca.p12")
michael@0: 
michael@0:     pwfile.write("\n")
michael@0: 
michael@0:     if not dbFilesExist(srcDir):
michael@0:       # Make sure all DB files from src are really deleted
michael@0:       unlinkDbFiles(srcDir)
michael@0: 
michael@0:       # Create certification database for ssltunnel
michael@0:       status = runUtil(certutil, ["-N", "-d", srcDir, "-f", pwfile.name])
michael@0:       if status:
michael@0:         return status
michael@0: 
michael@0:       status = runUtil(pk12util, ["-i", pgoCAPath, "-w", pwfile.name, "-d", srcDir, "-k", pwfile.name])
michael@0:       if status:
michael@0:         return status
michael@0: 
michael@0:     # Generate automatic certificate
michael@0:     locations = ServerLocations(os.path.join(build.topsrcdir,
michael@0:                                              "build", "pgo",
michael@0:                                              "server-locations.txt"))
michael@0:     iterator = iter(locations)
michael@0: 
michael@0:     # Skips the first entry, I don't know why: bug 879740
michael@0:     iterator.next()
michael@0: 
michael@0:     locationsParam = ""
michael@0:     firstLocation = ""
michael@0:     for loc in iterator:
michael@0:       if loc.scheme == "https" and "nocert" not in loc.options:
michael@0:         customCertOption = False
michael@0:         customCertRE = re.compile("^cert=(?:\w+)")
michael@0:         for option in loc.options:
michael@0:           match = customCertRE.match(option)
michael@0:           if match:
michael@0:             customCertOption = True
michael@0:             break
michael@0: 
michael@0:         if not customCertOption:
michael@0:           if len(locationsParam) > 0:
michael@0:             locationsParam += ","
michael@0:           locationsParam += loc.host
michael@0: 
michael@0:           if firstLocation == "":
michael@0:             firstLocation = loc.host
michael@0: 
michael@0:     if not firstLocation:
michael@0:       print "Nothing to generate, no automatic secure hosts specified"
michael@0:     else:
michael@0:       createRandomFile(rndfile)
michael@0: 
michael@0:       runUtil(certutil, ["-D", "-n", "pgo server certificate", "-d", srcDir, "-z", rndfile.name, "-f", pwfile.name])
michael@0:       # Ignore the result, the certificate may not be present when new database is being built
michael@0: 
michael@0:       status = runUtil(certutil, ["-S", "-s", "CN=%s" % firstLocation, "-t", "Pu,,", "-c", "pgo temporary ca", "-m", "2", "-8", locationsParam, "-v", "120", "-n", "pgo server certificate", "-d", srcDir, "-z", rndfile.name, "-f", pwfile.name])
michael@0:       if status:
michael@0:         return status
michael@0: 
michael@0:   return 0
michael@0: 
michael@0: if len(sys.argv) == 1:
michael@0:   print "Specify --gen-server or --gen-ca"
michael@0:   sys.exit(1)
michael@0: 
michael@0: build = MozbuildObject.from_environment()
michael@0: certdir = os.path.join(build.topsrcdir, "build", "pgo", "certs")
michael@0: if sys.argv[1] == "--gen-server":
michael@0:   certificateStatus = createSSLServerCertificate(build, certdir)
michael@0:   if certificateStatus:
michael@0:     print "TEST-UNEXPECTED-FAIL | SSL Server Certificate generation"
michael@0: 
michael@0:   sys.exit(certificateStatus)
michael@0: 
michael@0: if sys.argv[1] == "--gen-ca":
michael@0:   certificateStatus = createCertificateAuthority(build, certdir)
michael@0:   if certificateStatus:
michael@0:     print "TEST-UNEXPECTED-FAIL | Certificate Authority generation"
michael@0:   else:
michael@0:     print "\n\n"
michael@0:     print "==================================================="
michael@0:     print " IMPORTANT:"
michael@0:     print " To use this new certificate authority in tests"
michael@0:     print " run 'make' at testing/mochitest"
michael@0:     print "==================================================="
michael@0: 
michael@0:   sys.exit(certificateStatus)
michael@0: 
michael@0: print "Invalid option specified"
michael@0: sys.exit(1)