# Copyright (C) 2008 Todd Zullinger <tmz@pobox.com>
#
# 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 2 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/>.

"""Create a host-specific package for bootstrapping a puppet."""

version_info = (0, 6, 0)
version = '.'.join(map(str, version_info))

import os
import sys
import time
import shutil
import socket
import tarfile
import commands
import tempfile

defaults = {
    'destssldir': '',
    'domain': '',
    'force': False,
    'force_cert': False,
    'force_tarball': False,
    'force_package': False,
    'release': '1',
    'rpmdir': os.path.abspath(os.curdir),
    'sign': True,
    'ssldir': '/etc/puppet/ssl',
    'template': '%(ssldir)s/template.spec',
    'verbose': 1,
}

try:
    defaults['domain'] = socket.getfqdn().split('.', 1)[1]
except:
    pass

config = os.path.expanduser('~/.puppethost')
if os.path.exists(config):
    execfile(config, {}, defaults)
    truths = ['true', 'yes', '1']
    # Ensure boolean options are set correctly
    for opt in ['force', 'force_cert', 'force_tarball', 'force_package']:
        defaults[opt] = str(defaults[opt]).lower() in truths
    # Ensure numeric options are set correctly
    for opt in ['verbose']:
        if isinstance(defaults[opt], bool):
            defaults[opt] = int(defaults[opt])
        elif not isinstance(defaults[opt], int):
            try:
                defaults[opt] = int(str(defaults[opt]).lower() in truths)
            except:
                raise
                defaults[opt] = 0

package_types = ['deb', 'rpm']

def sign(packages, key='', verbose=defaults['verbose']):
    for type in package_types:
        pkgs = [p for t,p in packages if t == type]
        if not pkgs:
            continue
        try:
            sign_func = getattr(sys.modules[__name__], '_' + type + 'sign')
        except AttributeError:
            raise NotImplementedError(
                    'Signing of %s packages is not implemented' % type)
        sign_func(pkgs, key)

def _rpmsign(rpms, key='', tries=3):
    cmd = 'rpm --addsign'
    if key:
        cmd = cmd + ' --define "_gpg_name %s"' % key
    cmd += ' ' + ' '.join(rpms)
    while tries:
        tries -= 1
        status, output = commands.getstatusoutput(cmd)
        if status:
            if tries:
                print 'Error signing rpms:\n', output
                print 'Try again (%d left)' % tries
                continue
            raise PuppetHostError('Error signing rpms:\n%s' % output)
        else:
            print 'Sucessfully signed %d rpm(s)' % len(rpms)
            break
    return True

class PuppetHostError(StandardError):
    pass

class PuppetHost(object):
    def __init__(self, hostname = '', opts = defaults):
        """Puppet host class"""
        if not hostname:
            raise PuppetHostError('A hostname must be provided')

        if '.' not in hostname and opts['domain']:
            hostname = '.'.join([hostname, opts['domain']])
        hostname = hostname.lower()

        self.hostname = hostname
        self.opts = opts

        self.datefmt = '%Y%m%d'
        self.packages = []
        self.tarball = ''
        self._version = ''

    def gencert(self):
        """Generate a puppet certificate"""
        cert = self.files['cert']
        ssldir = self.opts['ssldir']
        if os.path.exists(cert) and not self.opts['force_cert']:
            raise PuppetHostError('%s exists, not overwriting' % cert)

        if self.opts['verbose']:
            print 'Generating certificate (%s)...' % cert,
            sys.stdout.flush()

        cmd = 'puppetca --generate --ssldir %s %s' % (ssldir, self.hostname)
        status, output = commands.getstatusoutput(cmd)
        if status:
            raise PuppetHostError('Error generating cert for %s: %s' %
                                    (self.hostname, output))

        if self.opts['verbose']:
            print 'done'

        self._version = time.strftime(self.datefmt)

    def package(self, packages = ['rpm']):
        """Create packages in requested formats"""
        for package in packages:
            if package not in package_types or not hasattr(self, package):
                raise PuppetHostError('Bogus package type: %s' % package)
            try: getattr(self, package)()
            except Exception, error:
                print error
                continue

    def deb(self):
        """Create a .deb package"""
        raise NotImplementedError('FIXME if you need debian packages.')

    def tar(self, path = ''):
        self._check_files()

        if not path:
            path = tempfile.mkdtemp('', 'puppet-host-')

        name = 'puppet-%s-%s' % (self.hostname, self.version)
        tarball = '%s/%s.tar.gz' % (path, name)

        if os.path.exists(tarball) and not self.opts['force_tarball']:
            raise PuppetHostError('%s exists, not overwriting' % tarball)

        if self.opts['verbose']:
            print 'Creating tarball (%s)...' % tarball,
            sys.stdout.flush()

        tar = tarfile.open(tarball, 'w:gz')
        tar.dereference = True
        for f in sorted(self.files.values()):
            arcname = '%s%s' % (name, f.replace(self.opts['ssldir'], ''))
            tar.add(f, arcname)
        tar.close()
        self.tarball = tarball

        if self.opts['verbose']:
            print 'done'

    def rpm(self):
        """Create a .rpm package"""
        rpmdir = self.opts['rpmdir']
        tmpdir = tempfile.mkdtemp('', 'puppet-host-')

        if not self.tarball or not os.path.exists(self.tarball):
            try: self.tar(tmpdir)
            except:
                self._cleanup(tmpdir)
                raise

        spec = '%s/puppet-%s.spec' % (tmpdir, self.hostname)

        text = open(self.opts['template']).read()
        text = text.replace('__HOSTNAME__', self.hostname)
        text = text.replace('__VERSION__', self.version)
        text = text.replace('__RELEASE__', self.opts['release'])
        text = text.replace('__SSLDIR__', self.opts['destssldir'])

        specfp = open(spec, 'w')
        specfp.write(text)
        specfp.close()

        rpm_fmt = '%{N}-%{V}-%{R}.%{ARCH}.rpm'
        cmd = 'rpm -q --qf "%s/%s" --specfile %s' % (rpmdir, rpm_fmt, spec)
        status, rpm = commands.getstatusoutput(cmd)
        if status:
            self._cleanup(tmpdir)
            raise PuppetHostError('Failed to get rpm filename:\n%s' % rpm)

        if os.path.exists(rpm) and not self.opts['force_package']:
            self._cleanup(tmpdir)
            raise PuppetHostError('%s exists, not overwriting' % rpm)

        if self.opts['verbose']:
            print 'Creating rpm (%s)...' % rpm,
            sys.stdout.flush()

        cmd = 'rpmbuild -bb'
        cmd += ' --define "_build_name_fmt %s"' % rpm_fmt
        for i in ['_builddir', '_sourcedir', '_specdir', '_srcrpmdir']:
            cmd += ' --define "%s %s"' % (i, tmpdir)
        cmd += ' --define "_rpmdir %s" %s' % (rpmdir, spec)

        status, output = commands.getstatusoutput(cmd)
        if status or not os.path.exists(rpm):
            self._cleanup(tmpdir)
            raise PuppetHostError('Error building rpm:\n', output)

        if self.opts['verbose']:
            print 'done'

        self._cleanup(tmpdir)
        self.packages.append(('rpm', rpm))

    def _check_files(self):
        """Check for the files we care about"""
        self.missing = [f for f in self.files.values()
                        if not os.path.exists(f)]
        if self.missing:
            error = 'The following required files are missing:\n'
            for f in self.missing:
                error += '\t%s' % f
            raise PuppetHostError(error)

    def _cleanup(self, dir):
        if os.path.isdir(dir):
            shutil.rmtree(dir)

    @property
    def files(self):
        hostname = self.hostname
        ssldir = self.opts['ssldir']
        cert = '%s/certs/%s.pem' % (ssldir, hostname)
        casigned = '%s/ca/signed/%s.pem' % (ssldir, hostname)
        if not os.path.exists(cert) and os.path.exists(casigned):
            os.symlink(casigned.replace(ssldir, '..'), cert)
        return {
            'crl':      '%s/ca/ca_crl.pem' % ssldir,
            'cacert':   '%s/certs/ca.pem' % ssldir,
            'cert':     '%s/certs/%s.pem' % (ssldir, hostname),
            'private':  '%s/private_keys/%s.pem' % (ssldir, hostname),
        }

    @property
    def version(self):
        """Get the date from a certificate"""
        if self._version:
            return self._version

        self._check_files()
        cert = self.files['cert']

        try:
            import OpenSSL.crypto as crypto
            x509 = crypto.load_certificate(crypto.FILETYPE_PEM,
                                           open(cert).read())
            self._version = x509.get_notBefore()[:8]
        except:
            # pyOpenSSL < 0.7 lacks load_certificate()... do it the hard way
            format = self.datefmt
            cmd = 'openssl x509 -in %s -text -noout' % cert
            status, output = commands.getstatusoutput(cmd)
            if status == 0:
                date = [l.split(':', 1)[1].lstrip() for l in output.split('\n')
                        if 'Not Before:' in l][0]
                fmt = '%b %d %H:%M:%S %Y %Z'
                self._version = time.strftime(format, time.strptime(date, fmt))
            else:
                self._version = time.strftime(format)
        return self._version