# Copyright 2011  Lars Wirzenius
# 
# 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/>.


'''System test classes.

This is like ``unittest``, except for testing running systems.

An extension is that the test runner will set up a member, ``settings``,
which acts as a dictionary, and has run-time settings for things like
the hostname of the target system.

'''


import logging
import os
import re
import subprocess
import tempfile
import unittest


class TestCase(unittest.TestCase):

    '''Base class for system tests.
    
    This extends ``unittest.TestCase`` with some things that are 
    relevant to system tests.
    
    '''

    def runcmd(self, argv, stdin='', *args, **kwargs):
        logging.debug('systest.TestCase.runcmd: argv=%s' % repr(argv))
        p = subprocess.Popen(argv, 
                             stdin=subprocess.PIPE, 
                             stdout=subprocess.PIPE, 
                             stderr=subprocess.PIPE)
        out, err = p.communicate(stdin)
        return p.returncode, out, err

    def hostcmd(self, argv, *args, **kwargs):
        '''Run external command on test host.'''
        returncode, out, err = self.runcmd(argv, *args, **kwargs)
        if returncode:
            msg = 'host command failed: %s\n%s' % (' '.join(argv), err)
            raise AssertionError(msg)
        return out

    def targetcmd(self, argv, *args, **kwargs):
        '''Run command on test target.'''
        
        self.setup_known_hosts()

        prefix = ['ssh', 
                  '-l', self.settings['user'], 
                  '-p', str(self.settings['ssh-port']),
                  '-o', 'UserKnownHostsFile %s' % self.known_hosts,
                 ]
        if self.settings['user-ssh-private-key']:
            prefix += ['-i', self.settings['user-ssh-private-key']]
        prefix += [self.settings['target'], '--']

        returncode, out, err = self.runcmd(prefix + argv, *args, **kwargs)
        if returncode:
            msg = 'target command failed: %s\n%s' % (' '.join(argv), err)
            raise AssertionError(msg)
        return out

    def setup_known_hosts(self):
        '''Setup a temporary known hosts file, and add target's key there.'''
        
        fd, self.known_hosts = tempfile.mkstemp()
        out = self.hostcmd(['ssh-keyscan', self.settings['target']])
        os.write(fd, out)
        os.close(fd)

    def assertMatches(self, pat, text, msg=None):
        self.assert_(re.match(pat, text),
                     msg=('pattern %s does not match %s' % (pat, text)) or msg)

    def find_open_ports(self):
        '''Return list of open ports.
        
        Each port is shown as '22/tcp', where '22' is the port number,
        and 'tcp' indicates the TCP protocol. See nmap documentation for
        more info on ports.
        
        '''

        out = self.hostcmd(['nmap', self.settings['target']])
        return [line.split()[0]
                 for line in out.splitlines()
                 if ' open ' in line]