safekeep-test

#!/usr/bin/python

import getopt, glob, os, os.path, re, shutil, socket, sys
import commands, random, time, tempfile, traceback
from commands import mkarg

ssh_id_file_keys = None
perform_snapshots = True
safekeep_args = ''
test_reps=2

class TestFailure(Exception):
    pass

def mkssh(cmd='ssh'):
    args = [cmd]
    ssh_id_file = ssh_id_file_keys[0]
    if ssh_id_file:
        args.append('-i')
        args.append(mkarg(ssh_id_file))
    return ' '.join(args)

def rpipe(cmd, user, host, desc, cmd2, user2, host2):
    if host:
        cmd = '%s %s@%s %s' % (mkssh(), user, host, mkarg(cmd))
    if host2:
        cmd2 = '%s %s@%s %s' % (mkssh(), user2, host2, mkarg(cmd2))
    print cmd, '|', cmd2
    pout = os.popen(cmd, 'r')
    pin = os.popen(cmd2, 'w')
    for line in pout:
        pin.write(line)
    st1 = pout.close()
    st2 = pin.close()
    if desc and st1:
        raise TestFailure('Failed to %s on host %s' % (desc, host))
    if desc and st2:
        raise TestFailure('Failed to %s on host %s' % (desc, host2))
    if st1: return st1
    return st2

def rcmd(cmd, user, host, desc=None):
    if host:
        cmd = '%s %s@%s %s' % (mkssh(), user, host, mkarg(cmd))
    print cmd
    st = os.system(cmd)
    if desc and st:
        raise TestFailure('Failed to %s on host %s' % (desc, host))
    return st

def rcmdin(cmd, content, user, host, desc=None):
    if host:
        cmd = '%s %s@%s %s' % (mkssh(), user, host, mkarg(cmd))
    print cmd
    pipe = os.popen(cmd, 'w')
    pipe.write(content)
    st = pipe.close()
    if st and desc:
        raise TestFailure('Failed to %s for %s@%s' % (desc, user, host))
    return st

def rcmdout(cmd, user, host, desc=None):
    if host:
        cmd = '%s %s@%s %s' % (mkssh(), user, host, mkarg(cmd))
    print cmd
    out = os.popen(cmd, 'r')
    try:
        return out.read()
    finally:
        st = out.close()
        if desc and st:
            raise TestFailure('Failed to %s on host %s' % (desc, host))
        if st:
            return None

def readfile(file, user=None, host=None):
    if host:
        cmd = 'cat %s' % file
        return rcmdout(cmd, user, host, 'read file %s' % file)
    else:
        fin = open(file)
        try:
            return fin.read()
        finally:
            fin.close()

def writefile(file, content, perm=None, mode='w', user=None, host=None):
    """
    Creates a file with the given name, content, and mode.
    The file may be passed in by name, as an open file object, 
    or as  file descriptor. If passed as a file object or file descriptor,
    it will be closed after the content will be written to it.
    """
    if host:
        if mode.startswith('a'):
            cmd = 'cat >> %s' % mkarg(file)
        else:
            cmd = 'cat > %s' % mkarg(file)
        rcmdin(cmd, content, user, host, 'write file %s' % file)
        if perm != None:
            cmd = 'chmod %s %s' % (perm, file)
            rcmd(cmd, user, host, 'set file permissions for ' + file)
    else:
        fout = open(file, mode)
        try:
            fout.write(content)
        finally:
            fout.close()

        if perm != None and isinstance(file, types.StringTypes): 
            os.chmod(file, perm)

def augmentfile(file, user=None, host=None):
    content = ""
    for i in range(random.randint(5, 20)):
        content = '%sSome random nr: %s\n' % (content, random.random())
    writefile(file, content, mode='a', user=user, host=host)

def localTest(tmproot):
    params = {'tmproot': tmproot, 'args': safekeep_args}
    os.mkdir(os.path.join(tmproot, 'backup.d'))
    os.mkdir(os.path.join(tmproot, 'client'))
    os.mkdir(os.path.join(tmproot, 'client', 'data'))
    os.mkdir(os.path.join(tmproot, 'client', 'home'))
    os.mkdir(os.path.join(tmproot, 'client', 'misc'))
    os.mkdir(os.path.join(tmproot, 'server'))
    CONFIG = """
        # Default values should do for now
    """
    writefile(os.path.join(tmproot, 'safekeep.conf'), CONFIG % params)
    BACKUP = """
       <backup id="%(tmproot)s">
          <repo path="%(tmproot)s/server" />
          <data>
             <exclude glob="**/*.nobackup" />
             <include path="%(tmproot)s/client" />
          </data>
        </backup>
    """
    writefile(os.path.join(tmproot, 'backup.d', 'test.backup'), BACKUP % params)
    for i in xrange(test_reps):
        FILES = (
            'data/fileA.out', 
            'data/fileB.nobackup', 
            'data/fileC.out', 
            'home/unit.out', 
            'home/office.out', 
            'home/street.out', 
            'misc/file.nobackup', 
        )
        for file in FILES:
            augmentfile(os.path.join(tmproot, 'client', file))
        cmd = "cd %(tmproot)s/client; find -type f -a \! -name '*.nobackup' | sort | xargs md5sum > %(tmproot)s/md5sums.client"
        if os.system(cmd % params):
            raise TestFailure("Can't compute the source MD5 sums")
        cmd = "safekeep --server %(args)s --conf '%(tmproot)s/safekeep.conf'"
        if os.system(cmd % params):
            raise TestFailure("Can't backup files")
        os.mkdir(os.path.join(tmproot, 'restore'))
        cmd = "rdiff-backup -r now %(tmproot)s/server %(tmproot)s/restore"
        if os.system(cmd % params):
            raise TestFailure("Can't restore files")
        cmd = "cd %(tmproot)s/restore%(tmproot)s/client; find -type f | sort | xargs md5sum > %(tmproot)s/md5sums.restore"
        if os.system(cmd % params):
            raise TestFailure("Can't compute the source MD5 sums")
        cmd = "diff -u %(tmproot)s/md5sums.client %(tmproot)s/md5sums.restore"
        if os.system(cmd % params):
            raise TestFailure("The MD5 sums differ")
        dir = os.path.join(tmproot, 'restore')
        os.system("chmod +w " + dir)
        shutil.rmtree(dir)
        time.sleep(1) 

def takeOver(host, role, ver):
    cmd = 'ssh -o PasswordAuthentication=no root@%s true' % (host)
    if os.system(cmd): 
        print 'The box %s does not appear to have been initialized' % (host)
        print 'for functioning as %s in a SafeKeep remote test.' % (role)
        print 'If you want to use it in such a role, please enter the root'
        print 'password for the box so that the test can take over it.'
        print 'WARNING: the test is potentially dangereous, all data on'
        print '         the box may be lost! Please make SURE this is'
        print '         the box you REALLY intend to use for the test.'
        print 'To cancel the test, type Ctrl-C'
        # add the key blindly to make sure we have it, and we are asked the passwd only once
        cmd = 'umask 077; test -d .ssh || mkdir .ssh; cat >> .ssh/authorized_keys'
        rcmdin(cmd, ssh_id_file_keys[1], 'root', host)

    cmd = 'cat /etc/redhat-release'
    release = rcmdout(cmd, 'root', host, 'fetch the release')
    res = re.match('\s*(.*) release (.*) \(.*\)\s*', release)
    if not res:
        raise TestFailure('Unknown release format for: %s' % release)
    distro = res.group(1)
    if distro in ('CentOS', ):
        distro = 'centos'
    elif distro in ('Fedora', 'Fedora Core'):
        distro = 'fedora'
    else:
        raise TestFailure('Unknown distro %s in release %s' % (distro, release))

    lattica_repo = """
        [safekeep-test-lattica-devel]
        name=Lattica - Development Tree for Fedora Core $releasever $basearch
        baseurl=http://lattica.com/repos/lattica/devel/%s/$releasever
        gpgkey=http://lattica.com/keys/RPM-GPG-KEY-lattica-devel
        enabled=1
        gpgcheck=0
    """ % (distro)
    lattica_repo = '\n'.join([line.strip() for line in lattica_repo.splitlines()])
    cmd = 'echo %s > /etc/yum.repos.d/safekeep-test-lattica-devel.repo' % (mkarg(lattica_repo).strip())
    rcmd(cmd, 'root', host, 'install Lattica Repo')
    cmd = """
        rm -rf /var/cache/yum/safekeep-test-lattica-devel/; 
        if rpm -q safekeep-%(role)s-%(verrel)s; then 
            yum update -y safekeep-%(role)s-%(verrel)s; 
        else 
            yum install -y safekeep-%(role)s-%(verrel)s;
        fi;
        rpm -q safekeep-%(role)s | grep %(verrel)s
    """  % { 'role': role, 'verrel': ver + '-1' }
    rcmd(cmd, 'root', host, 'install safekeep')

def createKey(user, host, keyname, comment):
    key = '.ssh/%s' % keyname
    cmd = 'test -f %s || ssh-keygen -q -b 1024 -t dsa -N \'\' -C \'%s\' -f %s' % (key, comment, key)
    rcmd(cmd, user, host, 'create %s on server' % keyname)
    return key

#TODO: this is just old code, we need to change it
#      so that it builds the RPM on the host it will run on
def packageAndDeployRPM(tar, user, host):
    pkgroot='/tmp/safekeep-test-XXX'   #TODO: replace XXX with curr timestamp

    rpmmacros = """
%%packager          safekeep Tester Package Builder
%%_topdir           %s
%%_debug_package    %%{nil}

%%_signature         gpg
""" % (pkgroot)

    writefile(os.path.join(pkgroot, '.rpmmacros'), rpmmacros)

    # and the directories needed by rpmbuild
    for dir in ('BUILD', 'RPMS/noarch', 'SOURCES', 'SPECS', 'SRPMS'):
        os.makedirs(os.path.join(pkgroot, dir))

    cmd = '%s %s %s@%s:%s' % \
          (mkssh('scp'), tar, user, host, pkgroot)
    print cmd
    if os.system(cmd):
         raise TestFailure('Failed to copy %s to the RPM host %s' % (tar, host)) 

    cmd = 'HOME=%s rpmbuild -ta %s' % (pkgroot, mytar)
    print cmd
    if os.system(cmd):
        raise TestFailure('Failed to build the .rpm')

    ver = mytar[len('safekeep-'):-len('.tar.gz')]
    binrpm_list = glob.glob(os.path.join(pkgroot, 'RPMS/noarch', 'safekeep-*-' + ver + '-*.noarch.rpm'))

    rel = binrpm_list[0]
    rel = rel[rel.rindex('-')+1:rel.rindex('.')]
    rel = rel[:rel.rindex('.')]

    for binrpm in binrpm_list:
        if not os.path.isfile(binrpm):
            raise TestFailure('Failed to find binary rpm: %s' % binrpm)

def packageAndDeployDEB(tar, user, host):
    # TODO: we need to implement this
    pass

def packageAndUpload(user, rpmhost, debhost):
    cmd = 'make tar'
    if os.system(cmd):
        raise TestFailure('Failed to build the tar snapshot')
    files = os.listdir('.')
    mytar = None
    for file in files:
        if file.startswith('safekeep-') and file.endswith('.tar.gz'):
            if file > mytar:
                mytar = file # last one wins
    if not mytar:
        raise TestFailure('Failed to determine the tar name')

    if rpmhost:
        pacakgeAndDeployRPM(mytar, user, rpmhost)

    if debhost:
        packageAndDeployDEB(mytar, user, debhost)

    cmd = 'rm %s' % (mytar)
    print cmd
    if os.system(cmd):
        raise TestFailure('Failed to nuke the tar')

    ver = mytar[len('safekeep-'):-len('.tar.gz')]
    return ver

def remoteTest(tmproot, client, server):
    # build and upload 
    ver = packageAndUpload('root', 'ulysses', None)

    takeOver(client, 'client', ver)
    takeOver(server, 'server', ver)

    # setup the server
    cmd = 'chsh -s /bin/bash safekeep --server'
    rcmd(cmd, 'root', server, 'Allow access to the safekeep account')
    client_addr = socket.gethostbyname(client)
    if client_addr != client:
        cmd = '(grep -iv %s /etc/hosts; echo "%s   %s") > /etc/hosts.new; mv -f /etc/hosts.new /etc/hosts' % (client, client_addr, client)
        rcmd(cmd, 'root', server, 'install the client name in /etc/hosts')

    cmd = 'cd ~safekeep; ' + \
          'cp /root/.ssh/authorized_keys .ssh/authorized_keys; ' + \
          'chown -R safekeep.safekeep .ssh'
    rcmd(cmd, 'root', server, 'install key for the safekeep user')
    key_id   = createKey('safekeep', server, 'id_dsa', 'SafeKeep server ID')
    key_ctrl = createKey('safekeep', server, 'safekeep-server-ctrl-key', 'SafeKeep server control key')
    key_data = createKey('safekeep', server, 'id-safekeep-data-key', 'SafeKeep server data key')

    cmd = 'if ! rpm -q postgresql-server; then p=postgresql-server; fi; ' + \
          'if ! rpm -q mysql-server; then p="$p mysql-server"; fi; ' + \
          'if [ "x$p" != "x" ]; then yum install -y $p; fi; ' + \
          'if service postgresql status | grep -v running; then service postgresql start; fi; ' + \
          'if service mysqld status | grep -v running; then service mysqld start; fi; '
    rcmd(cmd, 'root', server, 'ensure RDBMSes are installed and running')
   
    known_hosts = os.path.join(os.environ['HOME'], '.ssh', 'known_hosts')
    cmd = 'if test -f %s; then grep -i %s %s; fi' % (known_hosts, client, known_hosts)
    fingerprint = rcmdout(cmd, None, None, 'read the client fingerprint')
    if fingerprint:
        cmd = '(if test -f .ssh/known_hosts; then grep -iv %s .ssh/known_hosts; fi; cat) > .ssh/known_hosts.new; mv -f .ssh/known_hosts.new .ssh/known_hosts' % (client, )
        rcmdin(cmd, fingerprint, 'safekeep', server, 'deploy client fingerprint on server')

    snap_conf = ''
    if perform_snapshots:
        snap_conf = '<snapshot device="/dev/mapper/VolGroup00-LogVol00" size="50M"/>'
    conf = """
        <backup id="test-client">
          <host name="%s" key-data="%s" />
          <repo path="/var/lib/safekeep/client/data" retention="5h"/>
          <setup>
              %s
          </setup>
          <data>
             <exclude glob="**/*.nobackup" />
             <exclude path="/etc/cups/certs" />
             <exclude path="/etc/lvm" />
             <exclude path="/etc/mtab" />
             <include path="/etc" />
             <include path="/srv" />
          </data>
        </backup>
        """ % (client, key_data, snap_conf)

    writefile('/etc/safekeep/backup.d/test-client.backup', conf, '664', 'w', 'root', server)
    cmd = 'rm -rf client; mkdir -p client/data'
    rcmd(cmd, 'safekeep', server, 'create data repo') 

    # setup the client
    cmd = 'umask 077; test -d .ssh || mkdir .ssh; if test -f .ssh/authorized_keys; then cat .ssh/authorized_keys; fi'
    clkeys_txt = rcmdout(cmd, 'root', client, 'fetch client keys')
    key_id_enc = readfile('.ssh/id_dsa.pub', 'safekeep', server)
    if key_id_enc not in clkeys_txt.splitlines():
        writefile('.ssh/authorized_keys', key_id_enc, None, 'a', 'root', client)

    cmd = 'safekeep %s --keys --deploy' % safekeep_args
    rcmd(cmd, 'safekeep', server, 'deploy keys to client')

    # run the test
    for i in range(test_reps):
        cmd = 'rm -f md5sum.client md5sum.server'
        rcmd(cmd, 'safekeep', server, 'cleanup MD5 sums')
        FILES = (
            '/etc/a-simple-file', 
            '/etc/another-file.nobackup', 
            '/srv/safekeep-test1.log', 
            '/srv/safekeep-test2.nobackup',
        )
        for file in FILES:
            augmentfile(file, 'root', client)
        cmd1 = "cd /; (find /etc -type f; find /srv -type f) | " + \
               "grep -v '\.nobackup$' | grep -v '^/etc/cups/certs/' | " + \
               "grep -v '^/etc/mtab' | grep -v '^/etc/lvm/' | " + \
               "sort | sed s/.// | xargs md5sum"
        cmd2 = "cat > /var/lib/safekeep/md5sum.client"
        rpipe(cmd1, 'root', client, 'do MD5 sums on client', cmd2, 'safekeep', server)
        cmd = 'safekeep --server %s' % safekeep_args
        rcmd(cmd, 'root', server, 'backup data')
        cmd = 'cd /var/lib/safekeep/client/data; find -type f | grep -v rdiff-backup-data/ | sort | sed s/..// | xargs md5sum > /var/lib/safekeep/md5sum.server'
        rcmd(cmd, 'safekeep', server, 'do MD5 sums on server')
        cmd = 'diff -q /var/lib/safekeep/md5sum.client /var/lib/safekeep/md5sum.server'
        rcmd(cmd, 'safekeep', server, 'compare MD5 sums')
        
        # TODO: list available backups
        # TODO: compare against retention policy, fail if different
        cmd = 'if lvscan | grep _snap_; then false; fi'
        rcmd(cmd, 'root', client, 'test that no snapshots are left behind')
        cmd = 'if grep "safekeep-.*-rbind" /proc/mounts; then false; fi'
        rcmd(cmd, 'root', client, 'test that no --rbind mounts are left behind')
        # TODO: check iff appropriate DB dumps deleted
        time.sleep(1) 

def guess_ssh_id_file_keys():
    if os.environ.get('SSH_AUTH_SOCK'):
        out = os.popen('ssh-add -L', 'r')
        ids = out.read()
        if out.close() or not ids:
            print >> sys.stderr, 'Failed to get keys from ssh agent'
        else:
            return (None, ids)
    for filename in ('identity', 'id_dsa', 'id_rsa'):
        file = os.path.join(os.environ['HOME'], '.ssh', filename)
        if os.path.isfile(file): 
            return (file, readfile(file))
    return None

def usage():
    print 'usage: %s [options]' % (sys.argv[0])
    print
    print 'options:'
    print '-l,--local            Run the local version of the test (default)'
    print '-r,--remote           Run the client/server version of the test'
    print '-n,--nocleanup        Do not erase the temporary root for the test'
    print '-i,--identity IDFILE  SSH identity file to use (remote test only)'
    print '-c,--client=ADDR      The client address (remote test only)'
    print '-s,--server=ADDR      The server address (remote test only)'
    print '--reps=N              The number of iterations for the test (default 2)'
    print '--args=ARGS           Extra arguments to pass to SafeKeep (default to none)'
    print '-h,--help             Print this help message and exit'

def main():
    global ssh_id_file_keys
    try:
        opts, args = getopt.getopt(sys.argv[1:], 
                                   'lrni:c:s:', 
                                   ['local', 'remote', 'nocleanup', 'identity', 'reps=', 'args=', 'client=', 'server='])
    except getopt.GetoptError:
        usage()
        sys.exit(2)

    client = os.environ.get('SAFEKEEP_TEST_CLIENT')
    server = os.environ.get('SAFEKEEP_TEST_SERVER')
    nocleanup = False
    mode = 'local'
    for o, a in opts:
        if o in ('-l', '--local'):
            mode = 'local'
        elif o in ('-r', '--remote'):
            mode = 'remote'
        elif o in ('-n', '--nocleanup'):
            nocleanup = True
        elif o in ('-i', '--identity'):
            ssh_id_file_keys = (a, readfile(a))
        elif o in ('-c', '--client'):
            client = a
        elif o in ('-s', '--server'):
            server = a
        elif o in ('--reps', ):
            global test_reps
            test_reps = int(a)
        elif o in ('--args', ):
            global safekeep_args
            safekeep_args = a
        elif o in ('-h', '--help'):
            usage()
            sys.exit()

    if not client: client = 'safekeep-test-client'
    if not server: server = 'safekeep-test-server'

    exitcode = 0
    tmproot = tempfile.mkdtemp()
    try:
        try:
            if mode == 'remote':
                if not ssh_id_file_keys:
                    ssh_id_file_keys = guess_ssh_id_file_keys()
                assert ssh_id_file_keys, "Can not determine the SSH keys"
                remoteTest(tmproot, client, server)
            else:
                localTest(tmproot)
        except TestFailure, tf:
            print tf
            exitcode = 1
    finally:
        if nocleanup:
            print tmproot
        else:
            os.system("chmod -R +w " +tmproot)
            shutil.rmtree(tmproot)

    sys.exit(exitcode)

if __name__ == '__main__':
    main()

# vim: et ts=8 sw=4 sts=4