# 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, 8, 1)
version = '.'.join(map(str, version_info))

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

defaults = {
    'destdir': os.path.abspath(os.curdir),
    'destssldir': '',
    'domain': '',
    'force': False,
    'force_cert': False,
    'force_tarball': False,
    'force_package': False,
    'release': '1',
    'sign': True,
    'ssldir': '/etc/puppet/ssl',
    'template': '%(ssldir)s/template.spec',
    'package_type': 'rpm',
    'pkgprefix': 'host-package-',
    'puppetuser': 'puppet',
    'verbose': 1,
}

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

configs = ['/etc/puppet/hostpackage.conf', os.path.expanduser('~/.puppethost')]
for config in configs:
    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:
                    defaults[opt] = 0

package_types = ['deb', 'rpm', 'tar']

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

def _rpmsign(rpms, key='', tries=3):
    """Run rpm --addsign."""
    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):
    """Exception raised on errors."""
    pass

class PuppetHost(object):
    """PuppetHost class."""
    def __init__(self, hostname = '', opts = defaults):
        """Initialize PuppetHost objects."""
        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.missing = []
        self.packages = []
        self.tarball = ''
        self.tmpdir = tempfile.mkdtemp('', self.opts['pkgprefix'])
        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, ptype = None):
        """Create a package in the requested format."""
        if ptype is None:
            ptype = self.opts['package_type']
        if ptype not in package_types or not hasattr(self, ptype):
            raise PuppetHostError('Bogus package type: %s' % ptype)
        try:
            getattr(self, ptype)()
        except Exception, error:
            print error

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

    def tar(self):
        """Create a gzip'ed .tar file."""
        self._check_files()

        name = '%s%s-%s' % (self.opts['pkgprefix'], self.hostname, self.version)
        destdir = self.opts['package_type'] == 'tar' and self.opts['destdir'] or self.tmpdir
        tarball = '%s/%s.tar.gz' % (destdir, 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
        # Add the ssldir tarball if tar is the package type
        if self.opts['package_type'] == 'tar':
            tar.add(self.opts['ssldir'], recursive=False)
        for f in sorted(self.files.values()):
            arcname = None
            # Adjust the path in the archive if tar is not the package type
            if self.opts['package_type'] != 'tar':
                arcname = '%s%s' % (name, f.replace(self.opts['ssldir'], ''))
            tar.add(f, arcname)
        tar.close()
        self.tarball = tarball

        if self.opts['package_type'] == 'tar':
            self.packages.append(('tar', tarball))

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

    def rpm(self):
        """Create a .rpm package."""
        rpmdir = self.opts['destdir']

        if not self.tarball or not os.path.exists(self.tarball):
            self.tar()

        spec = '%s/puppet-%s.spec' % (self.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'])
        text = text.replace('__PKGPREFIX__', self.opts['pkgprefix'])
        text = text.replace('__PUPPETUSER__', self.opts['puppetuser'])

        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:
            raise PuppetHostError('Failed to get rpm filename:\n%s' % rpm)

        if os.path.exists(rpm) and not self.opts['force_package']:
            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, self.tmpdir)
        cmd += ' --define "_rpmdir %s" %s' % (rpmdir, spec)

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

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

        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, dirs=None):
        """Remove cruft."""
        dirs = dirs is not None and dirs or []
        dirs.append(self.tmpdir)
        for d in dirs:
            if os.path.isdir(d):
                shutil.rmtree(d)

    @property
    def files(self):
        """Expanded file list."""
        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
            datefmt = 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(datefmt, time.strptime(date, fmt))
            else:
                self._version = time.strftime(datefmt)
        return self._version