# Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import logging
import os

from autotest_lib.client.bin import utils
from autotest_lib.client.common_lib import error
from autotest_lib.client.cros import cros_ui_test

# Since we parse lsof output in several places, these centralize the
# column numbering for finding things in lsof output.  For example:
# autotest 1915 root 3u IPv4 9221 0t0 TCP *:https (LISTEN)
_LSOF_COMMAND = 0
_LSOF_PID = 1
_LSOF_USER = 2
_LSOF_FD = 3
_LSOF_TYPE = 4
_LSOF_DEVICE = 5
# In certain cases, the size/offset column is empty, making it more
# reliable to locate the last couple columns by counting from the right.
_LSOF_SIZE_OFF = 6
_LSOF_NODE = -3
_LSOF_NAME = -2

# We do this as a UITest so that we include any daemons that
# might be spawned at login, in our test results.
class security_NetworkListeners(cros_ui_test.UITest):
    """Check the system against a whitelist of expected network-listeners."""
    version = 1

    def load_baseline(self):
        """Loads the baseline of expected listeners."""
        # Figure out path to baseline file, by looking up our own path
        bpath = os.path.abspath(__file__)
        bpath = os.path.join(os.path.dirname(bpath), 'baseline')
        bfile = open(bpath)
        baseline_data = bfile.read()
        baseline_set = set([])
        for line in baseline_data.splitlines():
            line = line.strip()
            if len(line) != 0 and line[0] != '#': # skip comments
                baseline_set.add(line)
        bfile.close()
        return baseline_set


    def remove_autotest_noise(self, lsof_lines):
        """
        Processes underneath 'autotest' in the process tree
        unfortunately can inherit open sockets created by
        autotest. That leads to crazy-looking test failures where
        e.g. "sed" and "bash" appear to be listening on ports
        80/443. So, this takes the output of lsof and returns a
        filtered subset of it, with autotest stuff removed.

        @param lsof_lines: a list of lines as output by the 'lsof' util.
        """
        # Compile a set of the listening sockets to ignore.
        sockets_to_ignore = set([])
        for line in lsof_lines:
            fields = line.split()
            if fields[_LSOF_COMMAND] == 'autotest':
                sockets_to_ignore.add(fields[_LSOF_DEVICE])
        # Now that we know which ones to ignore, iterate the output again.
        lines_to_keep = []
        for line in lsof_lines:
            fields = line.split()
            if fields[_LSOF_DEVICE] in sockets_to_ignore:
                logging.debug('Ignoring %s', line)
            else:
                lines_to_keep.append(line)
        return lines_to_keep


    def run_once(self):
        """
        Compare a list of processes, listening on TCP ports, to a
        baseline. Test fails if there are mismatches.
        """
        cmd = (r'lsof -n -i -sTCP:LISTEN | '
               # Workaround for crosbug.com/28235 using a dynamic port #.
               r'sed "s/\\(shill.*127.0.0.1\\):.*/\1:DYNAMIC LISTEN/g"')
        cmd_output = utils.system_output(cmd, ignore_status=True,
                                         retain_output=True)
        # Use the [1:] slice to discard line 0, the lsof output header.
        lsof_lines = cmd_output.splitlines()[1:]
        # Unlike ps, we don't have a format option so we have to parse
        # lines that look like this:
        # sshd 1915 root 3u IPv4 9221 0t0 TCP *:ssh (LISTEN)
        # Out of that, we just want e.g. sshd *:ssh
        observed_set = set([])
        for line in self.remove_autotest_noise(lsof_lines):
            fields = line.split()
            observed_set.add('%s %s' % (fields[_LSOF_COMMAND],
                                        fields[_LSOF_NAME]))

        baseline_set = self.load_baseline()

        # If something in the observed set is not
        # covered by the baseline...
        diff = observed_set.difference(baseline_set)
        if diff:
            for daemon in diff:
                logging.error('Unexpected network listener: %s', daemon)

        # Or, things in baseline are missing from the system:
        diff2 = baseline_set.difference(observed_set)
        if diff2:
            for daemon in diff2:
                logging.error('Missing expected network listener: %s', daemon)

        if diff or diff2:
            raise error.TestFail('Baseline mismatch')