#!/usr/bin/python2 -E
#
# Authors:
#   Jan Cholasta <jcholast@redhat.com>
#
# Copyright (C) 2013  Red Hat
# see file 'COPYING' for use and warranty information
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
# Prevent garbage from readline on standard output
# (see https://fedorahosted.org/freeipa/ticket/4064)
if not os.isatty(1):
    os.environ['TERM'] = 'dumb'
import sys
import syslog
import traceback
import tempfile
import shutil
import base64
import contextlib

from ipapython import ipautil
from ipapython.dn import DN
from ipalib import api, errors, pkcs10, x509
from ipaplatform.paths import paths
from ipaserver.plugins.ldap2 import ldap2
from ipaserver.install import cainstance

# This is a certmonger CA helper script for IPA CA subsystem cert renewal. See
# https://git.fedorahosted.org/cgit/certmonger.git/tree/doc/submit.txt for more
# info on certmonger CA helper scripts.

# Return codes. Names of the constants are taken from
# https://git.fedorahosted.org/cgit/certmonger.git/tree/src/submit-e.h
ISSUED = 0
WAIT = 1
REJECTED = 2
UNREACHABLE = 3
UNCONFIGURED = 4
WAIT_WITH_DELAY = 5
OPERATION_NOT_SUPPORTED_BY_HELPER = 6

@contextlib.contextmanager
def ldap_connect():
    conn = None
    try:
        conn = ldap2(shared_instance=False, ldap_uri=api.env.ldap_uri)
        conn.connect(ccache=os.environ['KRB5CCNAME'])
        yield conn
    finally:
        if conn is not None and conn.isconnected():
            conn.disconnect()

def request_cert():
    """
    Request certificate from IPA CA.
    """
    syslog.syslog(syslog.LOG_NOTICE,
                  "Forwarding request to dogtag-ipa-renew-agent")

    path = paths.DOGTAG_IPA_RENEW_AGENT_SUBMIT
    args = [path] + sys.argv[1:]
    stdout, stderr, rc = ipautil.run(args, raiseonerr=False, env=os.environ)
    sys.stderr.write(stderr)
    sys.stderr.flush()

    syslog.syslog(syslog.LOG_NOTICE, "dogtag-ipa-renew-agent returned %d" % rc)

    if stdout.endswith('\n'):
        stdout = stdout[:-1]

    if rc == WAIT_WITH_DELAY:
        delay, sep, cookie = stdout.partition('\n')
        return (rc, delay, cookie)
    else:
        return (rc, stdout)

def store_cert():
    """
    Store certificate in LDAP.
    """
    operation = os.environ.get('CERTMONGER_OPERATION')
    if operation == 'SUBMIT':
        attempts = 0
    elif operation == 'POLL':
        cookie = os.environ.get('CERTMONGER_CA_COOKIE')
        if not cookie:
            return (UNCONFIGURED, "Cookie not provided")

        try:
            attempts = int(cookie)
        except ValueError:
            return (UNCONFIGURED, "Invalid cookie: %r" % cookie)
    else:
        return (OPERATION_NOT_SUPPORTED_BY_HELPER,)

    csr = os.environ.get('CERTMONGER_CSR')
    if not csr:
        return (UNCONFIGURED, "Certificate request not provided")

    nickname = pkcs10.get_friendlyname(csr)
    if not nickname:
        return (REJECTED, "No friendly name in the certificate request")

    cert = os.environ.get('CERTMONGER_CERTIFICATE')
    if not cert:
        return (REJECTED, "New certificate requests not supported")

    dercert = x509.normalize_certificate(cert)

    dn = DN(('cn', nickname), ('cn', 'ca_renewal'),
            ('cn', 'ipa'), ('cn', 'etc'), api.env.basedn)
    try:
        with ldap_connect() as conn:
            try:
                entry = conn.get_entry(dn, ['usercertificate'])
                entry['usercertificate'] = [dercert]
                conn.update_entry(entry)
            except errors.NotFound:
                entry = conn.make_entry(
                    dn,
                    objectclass=['top', 'pkiuser', 'nscontainer'],
                    cn=[nickname],
                    usercertificate=[dercert])
                conn.add_entry(entry)
            except errors.EmptyModlist:
                pass
    except Exception, e:
        attempts += 1
        if attempts < 10:
            syslog.syslog(
                syslog.LOG_ERR,
                "Updating renewal certificate failed: %s. Sleeping 30s" % e)
            return (WAIT_WITH_DELAY, 30, attempts)
        else:
            syslog.syslog(
                syslog.LOG_ERR,
                "Giving up. To retry storing the certificate, resubmit the "
                "request with profile \"ipaStorage\"")

    return (ISSUED, cert)

def request_and_store_cert():
    """
    Request certificate from IPA CA and store it in LDAP.
    """
    operation = os.environ.get('CERTMONGER_OPERATION')
    if operation == 'SUBMIT':
        state = 'request'
        cookie = None
    elif operation == 'POLL':
        cookie = os.environ.get('CERTMONGER_CA_COOKIE')
        if not cookie:
            return (UNCONFIGURED, "Cookie not provided")

        state, sep, cookie = cookie.partition(':')
        if state not in ('request', 'store'):
            return (UNCONFIGURED,
                    "Invalid cookie: %r" % os.environ['CERTMONGER_CA_COOKIE'])
    else:
        return (OPERATION_NOT_SUPPORTED_BY_HELPER,)

    if state == 'request':
        if cookie is None:
            os.environ['CERTMONGER_OPERATION'] = 'SUBMIT'
        else:
            os.environ['CERTMONGER_CA_COOKIE'] = cookie

        result = request_cert()
        if result[0] == WAIT:
            return (result[0], 'request:%s' % result[1])
        elif result[0] == WAIT_WITH_DELAY:
            return (result[0], result[1], 'request:%s' % result[2])
        elif result[0] != ISSUED:
            return result
        else:
            cert = result[1]
            cookie = None
    else:
        cert, sep, cookie = cookie.partition(':')

    if cookie is None:
        os.environ['CERTMONGER_OPERATION'] = 'SUBMIT'
    else:
        os.environ['CERTMONGER_CA_COOKIE'] = cookie
    os.environ['CERTMONGER_CERTIFICATE'] = cert

    result = store_cert()
    if result[0] == WAIT:
        return (result[0], 'store:%s:%s' % (cert, result[1]))
    elif result[0] == WAIT_WITH_DELAY:
        return (result[0], result[1], 'store:%s:%s' % (cert, result[2]))
    else:
        return result

def retrieve_cert():
    """
    Retrieve new certificate from LDAP.
    """
    operation = os.environ.get('CERTMONGER_OPERATION')
    if operation == 'SUBMIT':
        attempts = 0
    elif operation == 'POLL':
        cookie = os.environ.get('CERTMONGER_CA_COOKIE')
        if not cookie:
            return (UNCONFIGURED, "Cookie not provided")

        try:
            attempts = int(cookie)
        except ValueError:
            return (UNCONFIGURED, "Invalid cookie: %r" % cookie)
    else:
        return (OPERATION_NOT_SUPPORTED_BY_HELPER,)

    csr = os.environ.get('CERTMONGER_CSR')
    if not csr:
        return (UNCONFIGURED, "Certificate request not provided")

    nickname = pkcs10.get_friendlyname(csr)
    if not nickname:
        return (REJECTED, "No friendly name in the certificate request")

    old_cert = os.environ.get('CERTMONGER_CERTIFICATE')
    if not old_cert:
        return (REJECTED, "New certificate requests not supported")
    old_cert = x509.normalize_certificate(old_cert)

    syslog.syslog(syslog.LOG_NOTICE, "Updating certificate for %s" % nickname)

    with ldap_connect() as conn:
        try:
            entry = conn.get_entry(
                DN(('cn', nickname), ('cn', 'ca_renewal'),
                   ('cn', 'ipa'), ('cn', 'etc'), api.env.basedn),
                ['usercertificate'])
        except errors.NotFound:
            cert = old_cert
        else:
            cert = entry.single_value['usercertificate']

        if cert == old_cert:
            attempts += 1
            if attempts < 4:
                syslog.syslog(
                    syslog.LOG_INFO,
                    "Updated certificate for %s not available" % nickname)
                # No cert available yet, tell certmonger to wait another 8 hours
                return (WAIT_WITH_DELAY, 8 * 60 * 60, attempts)

        cert = base64.b64encode(cert)
        cert = x509.make_pem(cert)

    return (ISSUED, cert)

def export_csr():
    """
    This does not actually renew the cert, it just writes the CSR provided
    by certmonger to /var/lib/ipa/ca.csr and returns the existing cert.
    """
    operation = os.environ.get('CERTMONGER_OPERATION')
    if operation != 'SUBMIT':
        return (OPERATION_NOT_SUPPORTED_BY_HELPER,)

    csr = os.environ.get('CERTMONGER_CSR')
    if not csr:
        return (UNCONFIGURED, "Certificate request not provided")

    cert = os.environ.get('CERTMONGER_CERTIFICATE')
    if not cert:
        return (REJECTED, "New certificate requests not supported")

    csr_file = paths.IPA_CA_CSR
    try:
        with open(csr_file, 'wb') as f:
            f.write(csr)
    except Exception, e:
        return (UNREACHABLE, "Failed to write %s: %s" % (csr_file, e))

    return (ISSUED, cert)

def renew_ca_cert():
    """
    This is used for automatic CA certificate renewal.
    """
    cert = os.environ.get('CERTMONGER_CERTIFICATE')
    if not cert:
        return (REJECTED, "New certificate requests not supported")
    is_self_signed = x509.is_self_signed(cert)

    operation = os.environ.get('CERTMONGER_OPERATION')
    if operation == 'SUBMIT':
        state = 'retrieve'

        if is_self_signed:
            ca = cainstance.CAInstance(host_name=api.env.host, ldapi=False)
            if ca.is_renewal_master():
                state = 'request'
    elif operation == 'POLL':
        cookie = os.environ.get('CERTMONGER_CA_COOKIE')
        if not cookie:
            return (UNCONFIGURED, "Cookie not provided")

        state, sep, cookie = cookie.partition(':')
        if state not in ('retrieve', 'request'):
            return (UNCONFIGURED,
                    "Invalid cookie: %r" % os.environ['CERTMONGER_CA_COOKIE'])

        os.environ['CERTMONGER_CA_COOKIE'] = cookie
    else:
        return (OPERATION_NOT_SUPPORTED_BY_HELPER,)

    if state == 'retrieve':
        result = retrieve_cert()
        if result[0] == WAIT_WITH_DELAY and not is_self_signed:
            syslog.syslog(syslog.LOG_ALERT,
                          "IPA CA certificate is about to expire, "
                          "use ipa-cacert-manage to renew it")
    elif state == 'request':
        os.environ['CERTMONGER_CA_PROFILE'] = 'caCACert'
        result = request_and_store_cert()

    if result[0] == WAIT:
        return (result[0], '%s:%s' % (state, result[1]))
    elif result[0] == WAIT_WITH_DELAY:
        return (result[0], result[1], '%s:%s' % (state, result[2]))
    else:
        return result

def main():
    handlers = {
        'ipaStorage':       store_cert,
        'ipaRetrieval':     retrieve_cert,
        'ipaCSRExport':     export_csr,
        'ipaCACertRenewal': renew_ca_cert,
    }

    api.bootstrap(context='renew')
    api.finalize()

    operation = os.environ.get('CERTMONGER_OPERATION')
    if operation not in ('SUBMIT', 'POLL'):
        return OPERATION_NOT_SUPPORTED_BY_HELPER

    tmpdir = tempfile.mkdtemp(prefix="tmp-")
    try:
        principal = str('host/%s@%s' % (api.env.host, api.env.realm))
        ipautil.kinit_hostprincipal(paths.KRB5_KEYTAB, tmpdir, principal)

        profile = os.environ.get('CERTMONGER_CA_PROFILE')
        if profile:
            handler = handlers.get(profile, request_and_store_cert)
        else:
            ca = cainstance.CAInstance(host_name=api.env.host, ldapi=False)
            if ca.is_renewal_master():
                handler = request_and_store_cert
            else:
                handler = retrieve_cert

        res = handler()
        for item in res[1:]:
            print item
        return res[0]
    finally:
        shutil.rmtree(tmpdir)

try:
    sys.exit(main())
except Exception, e:
    syslog.syslog(syslog.LOG_ERR, traceback.format_exc())
    print "Internal error"
    sys.exit(UNREACHABLE)