# Copyright 2024 Red Hat, Inc.
# All Rights Reserved.

#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
import time

import fixtures
from oslo_log import log
from scapy.all import ICMP
from scapy.all import rdpcap
from tempest import config
from tempest.lib import exceptions

from whitebox_neutron_tempest_plugin.common import utils

WB_CONF = config.CONF.whitebox_neutron_plugin_options
LOG = log.getLogger(__name__)


class TcpdumpCapture(fixtures.Fixture):
    capture_files = None
    processes = None

    def __init__(self, client, interfaces, filter_str='', extra_prefix=''):
        self.client = client
        self.interfaces = [ifc.strip() for ifc in interfaces.split(',')]
        self.filter_str = filter_str
        self.timeout = WB_CONF.capture_timeout
        self.cmd_prefix = "{} sudo timeout {}".format(
            extra_prefix, self.timeout)

    def _setUp(self):
        self.start()

    def start(self):
        if not self.capture_files:
            # mktemp needs to be executed with sudo - otherwise the later
            # tcpdump command (run with sudo) fails because the created temp
            # file cannot be written
            # This happens in RHEL9 because fs.protected_regular is enabled
            self.capture_files = []
            self.processes = []

            for interface in self.interfaces:
                process = self.client.open_session()
                capture_file = self.client.exec_command('sudo mktemp').rstrip()
                cmd = '{} tcpdump -s0 -Uni {} {} -w {}'.format(
                    self.cmd_prefix, interface,
                    self.filter_str, capture_file)
                self.capture_files.append(capture_file)
                LOG.debug('Executing command: {}'.format(cmd))
                process.exec_command(cmd)
                self.processes.append(process)
            self.addCleanup(self.cleanup)

    def stop(self):
        for process in (self.processes or []):
            process.close()
        self.processes = None

    def cleanup(self):
        self.stop()
        if self.capture_files:
            if utils.host_responds_to_ping(self.client.host):
                self.client.exec_command(
                    '{} rm -f '.format(self.cmd_prefix) + ' '.join(
                        self.capture_files))
            self.capture_files = None

    def is_empty(self):
        try:
            pcap = rdpcap(self._open_capture_file())
        except Exception as e:
            LOG.debug('Error reading pcap file: ', str(e))
            return True
        for record in pcap:
            return False
        return True

    def get_next_hop_mtu(self):
        pcap = rdpcap(self._open_capture_file())
        for record in pcap:
            if 'IP' in record and 'ICMP' in record:
                icmp = record[ICMP]
                # ICMP type 3 = Destionation Unreachable
                if icmp.type == 3:
                    return repr(icmp.nexthopmtu)
        return None

    def _open_capture_file(self):
        if not self.capture_files:
            raise ValueError('No capture files available')
        elif len(self.capture_files) == 1:
            merged_cap_file = self.capture_files[0]
        else:
            cap_file_candidates = []
            print_pcap_file_cmd = '{} tcpdump -r {} | wc -l'
            for cap_file in self.capture_files:
                if 0 < int(self.client.exec_command(
                        print_pcap_file_cmd.format(
                            self.cmd_prefix, cap_file)).rstrip()):
                    # cap files that are not empty
                    cap_file_candidates.append(cap_file)

            if not cap_file_candidates:
                # they are all empty
                merged_cap_file = self.capture_files[0]
            elif 1 == len(cap_file_candidates):
                merged_cap_file = cap_file_candidates[0]
            else:
                merged_cap_file = self.client.exec_command(
                    self.cmd_prefix + ' mktemp').rstrip()
                n_retries = 5
                for i in range(n_retries):
                    try:
                        self.client.exec_command(
                            '{} tcpslice -w {} {}'.format(
                                self.cmd_prefix, merged_cap_file,
                                ' '.join(cap_file_candidates)))
                    except exceptions.SSHExecCommandFailed as exc:
                        if i == (n_retries - 1):
                            raise exc
                        LOG.warn('tcpslice command failed - retrying...')
                        time.sleep(5)
                    else:
                        break

        ssh_channel = self.client.open_session()
        ssh_channel.exec_command(self.cmd_prefix + ' cat ' + merged_cap_file)
        self.addCleanup(ssh_channel.close)
        return ssh_channel.makefile()