#!/usr/bin/env python
#
# Copyrite (c) 2014 SecurityKISS Ltd (http://www.securitykiss.com)  
#
# The MIT License (MIT)
#
# Yes, Mr patent attorney, you have nothing to do here. Find a decent job instead. 
# Fight intellectual "property".
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import subprocess, os, re, sys, time

def usage():
    usage = "rfwgen - key and certificate generator for rfw\n\nusage: rfwgen <server_ip>"
    return usage


def validate_ip(ip):
     """Check if the IP address has correct format.
     return validated and trimmed IP address as string or False if not valid
     """
     if not ip:
         return False
     ip = ip.strip()
     m = re.match(r"^(\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})$", ip)
     if m:
         a1, a2, a3, a4 = int(m.group(1)), int(m.group(2)), int(m.group(3)), int(m.group(4))
         if a1<256 and a2<256 and a3<256 and a4<256:
             ip_canon = "{}.{}.{}.{}".format(a1, a2, a3, a4)
             return ip_canon
     return False


def call(lcmd):
    try:
        out = subprocess.check_output(lcmd, stderr=subprocess.STDOUT)
        return out
    except subprocess.CalledProcessError, e:
        print("Error code {} returned when called '{}'. Command output: '{}'".format(e.returncode, e.cmd, e.output))
        raise e


def step(title):
    print('\n{}\n========================================='.format(title))

def abort():
    print('Aborted')
    sys.exit(1)

if __name__ == '__main__':

    if len(sys.argv) != 2:
        print usage()
        sys.exit(-1)
    server_ip = validate_ip(sys.argv[1])
    if not server_ip:
        print usage()
        sys.exit(-1)

    #TODO set umask and/or individual file permissions
    os.umask(0077)

    server_dir = 'server_{}'.format(server_ip)
    ca_key = 'offline/ca.key'
    ca_crt = 'client/ca.crt'
    server_key = os.path.join(server_dir, 'server.key')
    server_crt = os.path.join(server_dir, 'server.crt')
    server_csr = os.path.join(server_dir, 'server.csr')

    # check openssl if works
    try:
        call(['openssl', 'version'])
    except OSError, e:
        print('Could not run openssl. Check if it is correctly installed.')
        abort()
   
    # The procedure below based on http://blog.didierstevens.com/2008/12/30/howto-make-your-own-cert-with-openssl/
     
    # generate a 4096-bit long RSA key and the root CA if they don't exist
    if not os.path.isfile(ca_key) or not os.path.isfile(ca_crt):
        step('Create rfw root CA')
        print('Generating root CA in {} and {}'.format(ca_key, ca_crt))
        try:
            os.mkdir('offline')
            os.mkdir('client')
        except OSError, e:
            # ignore the 'dir already exists' error 
            pass
        root_subj = "/C=IE/ST=Universe/L=Internet/O=rfw root CA {}".format(time.strftime("%Y-%m-%d %H:%M:%S"))
        out = call(['openssl', 'req', '-new', '-newkey', 'rsa:4096', '-days', '7305', '-nodes', '-x509', '-subj', root_subj, '-keyout', ca_key, '-out', ca_crt])
        print(out)
        if not os.path.isfile(ca_key):
            print('Could not find ca_key. Something went wrong when generating {}'.format(ca_key))
            abort()
        if not os.path.isfile(ca_crt):
            print('Could not find ca_crt. Something went wrong when generating {}'.format(ca_crt))
            abort()
    else:
        step('Locate rfw root CA')
        print('Using existing root CA in {} and {}'.format(ca_key, ca_crt))
    

    step('Check existing server files for {}'.format(server_ip))
    if os.path.isfile(server_key) or os.path.isfile(server_crt):
        overwrite = raw_input('Server key or certificate already exists in folder {}\nDo you want to overwrite (y/n)? [n]: '.format(server_dir))
        if overwrite == 'y':
            try:
                if os.path.isfile(server_key):
                    os.remove(server_key)
                if os.path.isfile(server_crt):
                    os.remove(server_crt)
                if os.path.isfile(server_csr):
                    os.remove(server_csr)
            except OSError, e:
                print(e.message)
                print('Could not remove existing files in {}.'.format(server_dir))
                abort()
        else:
            abort()
    else:
        print('Done')
        
    try:
        os.mkdir(server_dir)
    except OSError, e:
        # ignore the 'dir already exists' error 
        pass

    step('Create server private key')
    out = call(['openssl', 'genrsa', '-out', server_key, '2048'])
    print(out)
    if not os.path.isfile(server_key):
        print('Could not find server_key. Something went wrong when generating {}'.format(server_key))
        abort()
    else:
        print('Done')

    step('Create certificate request')
    crt_subj = "/C=IE/ST=Universe/L=Internet/O=Server {0}/CN={0}".format(server_ip)
    out = call(['openssl', 'req', '-new', '-subj', crt_subj, '-key', server_key, '-out', server_csr])
    print(out)
    if not os.path.isfile(server_csr):
        print('Could not find server_csr. Something went wrong when generating {}'.format(server_csr))
        abort()
    else:
        print('Done')

    step('Create certificate signed by root CA')
    out = call(['openssl', 'x509', '-req', '-days', '7305', '-in', server_csr, '-CA', ca_crt, '-CAkey', ca_key, '-set_serial', '01', '-out', server_crt])
    print(out)
    if not os.path.isfile(server_crt):
        print('Could not find server_crt. Something went wrong when generating {}'.format(server_crt))
        abort()
    else:
        print('Done')

    os.remove(server_csr)

    print('\nrfwgen completed')    
    





