#!/usr/bin/python
#
#    test-open-iscsi.py quality assurance test script for open-iscsi
#    Copyright (C) 2011 Canonical Ltd.
#    Author: Jamie Strandboge <jamie@canonical.com>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License version 3,
#    as published by the Free Software Foundation.
#
#    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/>.
#
# packages required for test to run:
# QRT-Packages: open-iscsi
# packages where more than one package can satisfy a runtime requirement:
# QRT-Alternates: 
# files and directories required for the test to run:
# QRT-Depends: 
# privilege required for the test to run (remove line if running as user is okay):
# QRT-Privilege: root

'''
    In general, this test should be run in a virtual machine (VM) or possibly
    a chroot and not on a production machine. While efforts are made to make
    these tests non-destructive, there is no guarantee this script will not
    alter the machine. You have been warned.

    How to run in a clean VM:
    $ sudo apt-get -y install python-unit <QRT-Packages> && sudo ./test-PKG.py -v'

    How to run in a clean schroot named 'lucid':
    $ schroot -c lucid -u root -- sh -c 'apt-get -y install python-unit <QRT-Packages> && ./test-PKG.py -v'


    NOTES:
    - currently only tested on Ubuntu 8.04
'''


from netifaces import gateways, AF_INET
import unittest, subprocess, sys, os, glob
import shutil
import testlib
import textwrap
from tempfile import mkdtemp
import time

# There are setup based on README.multipurpose-vm. Feel free to override.
remote_server = ''
username = 'ubuntu'
password = 'passwd'
username_in = 'ubuntu'
password_in = 'ubuntupasswd'
initiatorname = 'iqn.2009-10.com.example.hardy-multi:iscsi-01'

COLLECT_USER_DATA = """\
#cloud-config
bukket:
 - &get_resolved_status |
   if [ "${1:--}" != "-" ]; then
      exec >"$1" 2>&1
   fi
   if ! command -v systemd-resolve >/dev/null 2>&1; then
       echo "systemd-resolve: not available."
       exit
   fi
   if ! ( systemd-resolve --help | grep -q -- --status ); then
       echo "systemd-resolve: no --status."
       exit
   fi
   systemd-resolve --status --no-pager

runcmd:
 - [ mkdir, -p, /output ]
 - [ cp, /etc/resolv.conf, /output]
 - [ sh, -c, *get_resolved_status, --, /output/systemd-resolve-status.txt]
 - [ tar, -C, /output, -cf, /dev/disk/by-id/virtio-output-disk, . ]

power_state:
 mode: poweroff
 message: cloud-init finished. Shutting down.
 timeout: 60
"""

try:
    from private.qrt.OpenIscsi import PrivateOpenIscsiTest
except ImportError:
    class PrivateOpenIscsiTest(object):
        '''Empty class'''
    print >>sys.stdout, "Skipping private tests"

class OpenIscsiTest(testlib.TestlibCase, PrivateOpenIscsiTest):
    '''Test my thing.'''

    def setUp(self):
        '''Set up prior to each test_* function'''
        self.pidfile = "/var/run/iscsid.pid"
        self.exe = "/sbin/iscsid"
        self.daemon = testlib.TestDaemon(["service", "open-iscsi"])
        self.initiatorname_iscsi = '/etc/iscsi/initiatorname.iscsi'
        self.iscsid_conf = '/etc/iscsi/iscsid.conf'

    def tearDown(self):
        '''Clean up after each test_* function'''
        global remote_server
        global initiatorname

        # If remote server is setup, convert back to manual, logout, remove
        # testlib configs and restart (in that order)
        if remote_server != '':
            testlib.cmd(['iscsiadm', '-m', 'node', '--targetname', initiatorname, '-p', '%s:3260' % remote_server, '--op=update', '--name', 'node.startup', '--value=manual'])
            testlib.cmd(['iscsiadm', '-m', 'node', '--targetname', initiatorname, '-p', '%s:3260' % remote_server, '--op=update', '--name',  'node.conn[0].startup', '--value=manual'])
            testlib.cmd(['iscsiadm', '-m', 'node', '--targetname', initiatorname, '-p', '%s:3260' % remote_server, '--logout'])

        testlib.config_restore(self.initiatorname_iscsi)
        testlib.config_restore(self.iscsid_conf)
        self.daemon.restart()

    def test_discovery_with_server(self):
        '''Test iscsi_discovery to remote server'''
        global remote_server
        global username
        global password
        global username_in
        global password_in
        global initiatorname

        if remote_server == '':
            return self._skipped("--remote-server not specified")

        contents = '''
InitiatorName=%s
InitiatorAlias=ubuntu
''' % (initiatorname)
        testlib.config_replace(self.initiatorname_iscsi, contents, True)

        contents = '''
node.session.auth.authmethod = CHAP
node.session.auth.username = %s
node.session.auth.password = %s
node.session.auth.username_in = %s
node.session.auth.password_in = %s

discovery.sendtargets.auth.authmethod = CHAP
discovery.sendtargets.auth.username = %s
discovery.sendtargets.auth.password = %s
discovery.sendtargets.auth.username_in = %s
discovery.sendtargets.auth.password_in = %s
''' % (username, password, username_in, password_in, username, password, username_in, password_in)
        testlib.config_replace(self.iscsid_conf, contents, True)

        self.assertTrue(self.daemon.restart())
        time.sleep(2)
        self.assertTrue(self.daemon.status()[0])

        rc, report = testlib.cmd(["/sbin/iscsi_discovery", remote_server])
        expected = 0
        result = 'Got exit code %d, expected %d\n' % (rc, expected)
        self.assertEquals(expected, rc, result + report)
        for i in ['starting discovery to %s' % remote_server,
                  'Testing iser-login to target %s portal %s' % (initiatorname, remote_server),
                  'starting to test tcp-login to target %s portal %s' % (initiatorname, remote_server),
                  'discovered 1 targets at %s, connected to 1' % remote_server]:
            result = "Could not find '%s' in report:\n" % i
            self.assertTrue(i in report, result + report)

    def test_net_interface_handler_execute_bit(self):
        '''Test /lib/open-iscsi/net-interface-handler is executable.'''
        nih_path = '/lib/open-iscsi/net-interface-handler'
        self.assertTrue(os.access(nih_path, os.X_OK))

class CloudImageTest(testlib.TestlibCase, PrivateOpenIscsiTest):
    '''Test the cloud image can boot on iscsi root.

    See README-boot-test.md for more information.
    '''

    @classmethod
    def setUpClass(cls):
        reason = (
            "Skipped Cloud Image test on {arch}. whitelisted are: {whitelist}")
        whitelisted = os.environ.get('WHITELIST_ARCHES', 'amd64').split(",")
        curarch = testlib.manager.dpkg_arch
        if testlib.manager.dpkg_arch not in whitelisted:
            raise unittest.SkipTest(
                reason.format(arch=curarch, whitelist=whitelisted))

        here = os.path.dirname(os.path.abspath(__file__))
        os.environ['PATH'] = "%s:%s" % (here, os.environ['PATH'])
        release = os.environ.get(
            "ISCSI_TEST_RELEASE", testlib.ubuntu_release())
        image_d = os.path.join(here, '{}.d'.format(release))
        # Download MAAS ephemeral image.
        info = {'release': release,
                'image_d': image_d,
                'root_image': os.path.join(image_d, 'disk.img'),
                'kernel': os.path.join(image_d, 'kernel'),
                'initrd': os.path.join(image_d, 'initrd')}
        try:
            get_image(info['image_d'], release)
        except CalledProcessError as e:
            if e.return_code != 3:
                raise e
            raise unittest.SkipTest(
                "Cloud Image not available for release '%s'." % release)
        if os.environ.get("NO_PATCH_IMAGE", "0") == "0":
            patch_image(info['root_image'],
                        kernel=info['kernel'], initrd=info['initrd'])
        cls.info = info
        cls.here = here

    def tearDown(self):
        super(CloudImageTest, self).tearDown()
        if os.path.exists(self.tmpdir):
            shutil.rmtree(self.tmpdir)

    def setUp(self):
        super(CloudImageTest, self).setUp()
        self.tmpdir = mkdtemp()
        udp = os.path.join(self.tmpdir, 'user-data')
        with open(udp, "w") as fp:
            fp.write(COLLECT_USER_DATA)
        self.info['user-data'] = udp

    def create_output_disk(self):
        path = os.path.join(self.tmpdir, 'output-disk.img')
        subprocess.check_call([
            'qemu-img', 'create', '-f',  'raw', path, '10M'])
	return path

    def extract_files(self, path):
        # get contents in a dictionary of tarball at 'path'
        tmpdir = mkdtemp()
        try:
            subprocess.check_call(['tar', '-C', tmpdir, '-xf', path])
            flist = {}
            prefix = len(tmpdir) + 1
            for root, dirs, files in os.walk(tmpdir):
                for fname in files:
                    fpath = os.path.join(root, fname)
                    key = fpath[prefix:]
                    with open(fpath, "r") as fp:
                        flist[key] = fp.read()
            return flist
        finally:
            shutil.rmtree(tmpdir)

    def test_tgt_boot(self):
        tgt_boot_cmd = os.path.join(self.here, 'tgt-boot-test')
        # Add self.here to PATH so xkvm will be available to tgt-boot-test
        dns_addr = '10.1.1.4'
        rel_dir = '{}.d'.format(self.info['release'])
        dns_search = 'example.com'
        info = {'host': '10.1.1.2', 'dns': dns_addr,
                'dnssearch': dns_search, 'network': '10.1.1.0/24'}
        netdev = ("--netdev=user,net={network},host={host},dns={dns},"
                  "dnssearch={dnssearch}").format(**info)

        output_disk = self.create_output_disk()
        cmd = [
            tgt_boot_cmd, '-v', netdev,
            '--disk=%s,serial=output-disk' % output_disk,
            '--user-data-add=%s' % self.info['user-data'],
            self.info['root_image'], self.info['kernel'],
            self.info['initrd']]
        sys.stderr.write(' '.join(cmd) + "\n")

        subprocess.check_call(cmd)
        files = self.extract_files(output_disk)
        print("files: %s" % files)
        resolvconf = files.get('resolv.conf', "NO_RESOLVCONF_FOUND")
        resolve_status = files.get('systemd-resolve-status.txt')

        resolvconf_id = 'generated by resolvconf'
        resolved_addr = "127.0.0.53"
        if resolvconf_id in resolvconf:
            print("resolvconf manages resolvconf.\n")
            self.assertIn(
                dns_addr, resolvconf,
                msg = ("%s not in resolvconf contents: \n%s" %
                       (dns_addr, resolvconf)))
            if dns_search in resolvconf:
                print("%s was found in resolv.conf." % dns_search)
            elif resolved_addr in resolvconf and dns_search in resolve_status:
                # zesty has resolvconf and systemd-resolved.
                print("%s was in resolve_status and %s in resolv.conf" %
                      (resolved_addr, dns_search))
            else:
                raise AssertionError(
                    "%s domain is not being searched." % dns_search)

        else:
            print("systemd-resolved managing resolve.conf\n")
            self.assertIn(
                resolved_addr, resolvconf,
                msg="%s not in resolved resolv.conf: %s" % (resolved_addr,
                                                            resolvconf))
            self.assertIn(dns_addr, resolve_status,
                msg=("%s not in systemd-resolve status: %s" %
                     (dns_addr, resolve_status)))
            self.assertIn(dns_search, resolve_status,
                msg=("search addr '%s' not in systemd-resolve status: %s" %
                     (dns_search, resolve_status)))


def get_image(top_d, release):
    cmd = ['get-image', top_d, 'cloud-daily', release]
    subprocess.check_call(cmd)


def patch_image(image, packages=None, kernel=None, initrd=None):
    '''Patch root-image with dep8 built open-iscsi package.'''

    if packages is None:
        # an empty 'packages' to patch-image still installs open-iscsi
        packages = []

    cmd = ['patch-image', image]
    if kernel:
        cmd.append('--kernel=%s' % kernel)
    if initrd:
        cmd.append('--initrd=%s' % initrd)
    cmd.extend(packages)

    subprocess.check_call(cmd)


if __name__ == '__main__':
    import optparse
    parser = optparse.OptionParser()
    parser.add_option("-v", "--verbose", dest="verbose", help="Verbose", action="store_true")
    parser.add_option("-s", "--remote-server", dest="remote_server", help="Specify host with available iSCSI volumes", metavar="HOST")

    parser.add_option("-n", "--initiatorname", dest="initiatorname", help="Specify initiatorname for use with --remote-server", metavar="NAME")

    parser.add_option("--password", dest="password", help="Specify password for use with --remote-server", metavar="PASS")
    parser.add_option("--password-in", dest="password_in", help="Specify password_in for use with --remote-server", metavar="PASS")

    parser.add_option("--username", dest="username", help="Specify username for use with --remote-server", metavar="USER")
    parser.add_option("--username-in", dest="username_in", help="Specify username_in for use with --remote-server", metavar="USER")

    (options, args) = parser.parse_args()

    if options.remote_server:
        remote_server = options.remote_server

        if options.username:
            username = options.username
        if options.password:
            password = options.password
        if options.username_in:
            username_in = options.username_in
        if options.password_in:
            password_in = options.password_in
        if options.initiatorname:
            initiatorname = options.initiatorname
        print "Connecting to remote server with:"
        print " host = %s " % remote_server
        print ' initiatorname = %s' % initiatorname
        print ' username = %s' % username
        print ' password = %s' % password
        print ' username_in = %s' % username_in
        print ' password_in = %s' % password_in

    suite = unittest.TestSuite()
    suite.addTest(unittest.TestLoader().loadTestsFromTestCase(OpenIscsiTest))
    suite.addTest(unittest.TestLoader().loadTestsFromTestCase(
        CloudImageTest))
    rc = unittest.TextTestRunner(verbosity=2).run(suite)
    if not rc.wasSuccessful():
        sys.exit(1)
