# Copyright 2012 Google Inc. All Rights Reserved.
# Modifications Copyright 2022-2023 Aerleon Project Authors.
#
# 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.
#

"""PacketFilter (PF) generator."""

import collections
import copy
from typing import Optional, Union, cast

from absl import logging

from aerleon.lib import aclgenerator, nacaddr
from aerleon.lib.nacaddr import IPv4, IPv6
from aerleon.lib.policy import Policy, Term, VarType


class Error(aclgenerator.Error):
    """Base error class."""


class DuplicateTermError(Error):
    """Raised when duplication of term names are detected."""


class DuplicateShortenedTableNameError(Error):
    """Raised when a duplicate shortened table name is found."""


class UnsupportedProtoError(Error):
    """Raised when a protocol is not supported."""


class Term(aclgenerator.Term):
    """Generate PacketFilter policy terms."""

    # Validate that term does not contain any fields we do not
    # support.  This prevents us from thinking that our output is
    # correct in cases where we've omitted fields from term.
    _PLATFORM = 'packetfilter'
    _ACTION_TABLE = {
        'accept': 'pass',
        'deny': 'block drop',
        'reject': 'block return',
        'next': 'pass',
    }
    # Moving the log keyword into an member variable allows subclasses to override
    # it to support logging options outside of the scope of the Aerleon policy
    # spec, e.g., on platform-specific options such as packetfilter's
    # "log (all, to pflogN)" per-direction.
    _LOG_TABLE = {
        '': 'log',
        'in': 'log',
        'out': 'log',
    }
    _QUICK_TABLE = {
        'accept': 'quick',
        'deny': 'quick',
        'reject': 'quick',
        'next': '',
    }
    _DIRECTION_TABLE = {
        '': '',
        'in': 'in',
        'out': 'out',
    }
    _TCP_FLAGS_TABLE = {
        'syn': 'S',
        'ack': 'A',
        'fin': 'F',
        'rst': 'R',
        'urg': 'U',
        'psh': 'P',
        'all': 'SAFRUP',
        'none': 'NONE',
    }
    _PROTO_TABLE = {
        'icmpv6': 'ipv6-icmp',
    }
    _UNSUPPORTED_PROTOS = ['hopopt']

    def __init__(
        self,
        term: Term,
        filter_name: str,
        stateful: bool = True,
        af: str = 'inet',
        direction: str = '',
    ) -> None:
        """Setup a new term.

        Args:
          term: A policy.Term object to represent in packetfilter.
          filter_name: The name of the filter chan to attach the term to.
          stateful: Whether to keep firewall state for the term.
          af: Which address family ('inet' or 'inet6') to apply the term to.
          direction: What direction the term applies to ('in', 'out' or both).

        Raises:
          aclgenerator.UnsupportedFilterError: Filter is not supported.
        """
        super().__init__(term)
        self.term = term  # term object
        self.filter = filter_name  # actual name of filter
        self.options = []
        self.default_action = 'deny'
        self.af = af
        self.stateful = stateful
        self.direction = direction

    def __str__(self) -> str:
        """Render config output from this term object."""

        ret_str = []
        self._SetDefaultAction()

        # Create a new term
        ret_str.append(f'\n# term {self.term.name}')

        comments = aclgenerator.WrapWords(self.term.comment, 80)
        # append comments to output
        if comments and comments[0]:
            for line in comments:
                ret_str.append(f'# {line!s}')

        if str(self.term.action[0]) not in self._ACTION_TABLE:
            raise aclgenerator.UnsupportedFilterError(
                '%s %s %s %s'
                % ('\n', self.term.name, self.term.action[0], 'action not currently supported.')
            )

        if self.direction and str(self.direction) not in self._DIRECTION_TABLE:
            raise aclgenerator.UnsupportedFilterError(
                '%s %s %s %s'
                % ('\n', self.term.name, self.term.direction, 'direction not currently supported.')
            )
        # protocol
        if self.term.protocol:
            protocol = self.term.protocol
        else:
            protocol = []

        # source address
        term_saddrs = self._CheckAddressAf(self.term.source_address)
        if not term_saddrs:
            if self.af != 'mixed':
                logging.warning(
                    self.NO_AF_LOG_ADDR.substitute(
                        term=self.term.name, direction='source', af=self.af
                    )
                )
            return ''
        term_saddr = self._GenerateAddrStatement(term_saddrs, self.term.source_address_exclude)

        # destination address
        term_daddrs = self._CheckAddressAf(self.term.destination_address)
        if not term_daddrs:
            if self.af != 'mixed':
                logging.warning(
                    self.NO_AF_LOG_ADDR.substitute(
                        term=self.term.name, direction='destination', af=self.af
                    )
                )
            return ''
        term_daddr = self._GenerateAddrStatement(
            term_daddrs, self.term.destination_address_exclude
        )

        # ports
        source_port = []
        destination_port = []
        if self.term.source_port:
            source_port = self._GeneratePortStatement(self.term.source_port)
        if self.term.destination_port:
            destination_port = self._GeneratePortStatement(self.term.destination_port)

        # icmp-type
        icmp_types = ['']
        if self.term.icmp_type:
            if protocol == ['icmp']:
                af = 'inet'
            elif protocol == ['icmpv6']:
                af = 'inet6'
            else:
                raise aclgenerator.UnsupportedFilterError(
                    '%s %s %s'
                    % ('\n', self.term.name, 'icmp protocol is not defined or not supported.')
                )
            icmp_types = self.NormalizeIcmpTypes(self.term.icmp_type, protocol, af)

        # options
        tcp_flags_set = []
        tcp_flags_check = []
        for next_opt in [str(x) for x in self.term.option]:
            for next_flag in self._TCP_FLAGS_TABLE:
                if next_opt.find(next_flag) == 0:
                    if protocol != ['tcp']:
                        raise aclgenerator.UnsupportedFilterError(
                            '%s %s %s'
                            % (
                                '\n',
                                self.term.name,
                                'tcp flags may only be specified with tcp protocol.',
                            )
                        )
                    tcp_flags_set.append(self._TCP_FLAGS_TABLE.get(next_flag))
                    tcp_flags_check.append(self._TCP_FLAGS_TABLE.get(next_flag))

        # If tcp-established is set, override any of the flags above with the
        # S/SA flags.  Issue an error if flags are specified with 'established'.
        for opt in [str(x) for x in self.term.option]:
            if opt == 'established' or opt == 'tcp-established':
                if tcp_flags_set or tcp_flags_check:
                    raise aclgenerator.UnsupportedFilterError(
                        '%s %s %s'
                        % (
                            '\n',
                            self.term.name,
                            'tcp flags may not be specified with tcp-established.',
                        )
                    )
                # We need to set 'flags A/A' for established regardless of whether or
                # not we're stateful:
                # - if we stateful, the default is 'flags S/SA' which prevent writing
                # rules for reply packets.
                # - if we're stateless, this is the only way to do it.
                if not protocol or 'tcp' in protocol:
                    tcp_flags_set.append(self._TCP_FLAGS_TABLE.get('ack'))
                    tcp_flags_check.append(self._TCP_FLAGS_TABLE.get('ack'))

        # The default behavior of pf is 'keep state flags S/SA'.  If we're not
        # stateless, and if flags have not been specified explicitly via options,
        # append that here.  Note that pf allows appending flags for udp and icmp;
        # they are just ignored, as long as TCP is in the proto.  This lets you
        # doing things like 'proto { tcp udp icmp } flags S/SA' and have the flags
        # only applied to the tcp bits that match.  However, the policy description
        # language prohibits setting flags on non-TCP, since it doesn't make sense
        # on all platforms.
        if (
            (not protocol or protocol == ['tcp'])
            and self.stateful
            and not tcp_flags_set
            and not tcp_flags_check
        ):
            tcp_flags_set.append(self._TCP_FLAGS_TABLE.get('syn'))
            tcp_flags_check.append(self._TCP_FLAGS_TABLE.get('syn'))
            tcp_flags_check.append(self._TCP_FLAGS_TABLE.get('ack'))

        ret_str.extend(
            self._FormatPart(
                self.term.action[0],
                self.direction,
                self.term.logging,
                self.term.interface,
                self.af,
                protocol,
                term_saddr,
                source_port,
                term_daddr,
                destination_port,
                tcp_flags_set,
                tcp_flags_check,
                icmp_types,
                self.options,
                self.stateful,
            )
        )

        return '\n'.join(str(v) for v in ret_str if v)

    def _CheckAddressAf(self, addrs: list[Union[IPv4, IPv6]]) -> list[Union[IPv4, IPv6, str]]:
        """Verify that the requested address-family matches the address's family."""
        if not addrs:
            return ['any']
        if self.af == 'mixed':
            return addrs
        af_addrs = []
        af = self.NormalizeAddressFamily(self.af)
        for addr in addrs:
            if addr.version == af:
                af_addrs.append(addr)
        return af_addrs

    def _FormatPart(
        self,
        action: str,
        direction: str,
        log: list[VarType],
        interface: Optional[str],
        af: str,
        proto: list[str],
        src_addr: str,
        src_port: Union[list[str], str],
        dst_addr: str,
        dst_port: Union[list[str], str],
        tcp_flags_set: list[str],
        tcp_flags_check: list[str],
        icmp_types: list[Union[str, int]],
        options: list[str],
        stateful: bool,
    ) -> list[str]:
        """Format the string which will become a single PF entry."""
        line = [f'{self._ACTION_TABLE.get(action)}']

        if direction:
            line.append(direction)

        quick = self._QUICK_TABLE.get(action)
        if quick:
            line.append(f'{quick}')

        if log:
            logaction = self._LOG_TABLE.get(direction)
            if logaction:
                line.append(logaction)
            else:
                line.append('log')

        if interface:
            line.append(f'on {interface}')

        if af != 'mixed':
            line.append(af)

        if proto:
            line.append(self._GenerateProtoStatement(proto))

        line.append(f'from {src_addr}')
        if src_port:
            line.append(f'port {src_port}')

        line.append(f'to {dst_addr}')
        if dst_port:
            line.append(f'port {dst_port}')

        if tcp_flags_set and tcp_flags_check:
            line.append('flags')
            line.append(f"{''.join(tcp_flags_set)}/{''.join(tcp_flags_check)}")

        if 'icmp' in proto and icmp_types:
            type_strs = [str(icmp_type) for icmp_type in icmp_types]
            type_strs = ', '.join(type_strs)
            if type_strs:
                line.append('icmp-type { %s }' % type_strs)

        if 'icmpv6' in proto and icmp_types:
            type_strs = [str(icmp_type) for icmp_type in icmp_types]
            type_strs = ', '.join(type_strs)
            if type_strs:
                line.append('icmp6-type { %s }' % type_strs)

        if options:
            line.extend(options)

        if not stateful:
            line.append('no state')
        elif action in ['accept', 'next']:
            line.append('keep state')

        return [' '.join(line)]

    def _GenerateProtoStatement(self, protocols: list[str]) -> str:
        proto = ''
        if protocols:
            protocols = copy.deepcopy(protocols)
            for i, proto in enumerate(protocols):
                if proto in self._UNSUPPORTED_PROTOS:
                    raise UnsupportedProtoError
                try:
                    protocols[i] = self._PROTO_TABLE[proto]
                except KeyError:
                    pass
            proto = 'proto { %s }' % ' '.join(protocols)
        return proto

    def _GenerateAddrStatement(
        self, addrs: list[Union[IPv4, IPv6]], exclude_addrs: list[Union[IPv4, IPv6]]
    ) -> str:
        addresses = set()
        if addrs != ['any']:
            parent_token_set = set()
            for addr in addrs:
                addr = cast(nacaddr.IPType, addr)
                parent_token_set.add(addr.parent_token)
            for token in parent_token_set:
                addresses.add(f'<{token[:31]}>')
        else:
            addresses.add('any')
        if exclude_addrs != ['any']:
            parent_token_set = set()
            for addr in exclude_addrs:
                addr = cast(nacaddr.IPType, addr)
                parent_token_set.add(addr.parent_token)
            for token in parent_token_set:
                addresses.add(f'!<{token[:31]}>')
        return '{ %s }' % ', '.join(sorted(addresses))

    def _GeneratePortStatement(self, ports: list[tuple[int, int]]) -> str:
        port_list = []
        for port_tuple in ports:
            if port_tuple[0] == port_tuple[1]:
                port_list.append(str(port_tuple[0]))
            else:
                port_list.append(f'{port_tuple[0]}:{port_tuple[1]}')
        return '{ %s }' % (' '.join(list(collections.OrderedDict.fromkeys(port_list))))

    def _SetDefaultAction(self) -> None:
        """If term does not specify action, use filter default action."""
        if not self.term.action:
            self.term.action[0].value = self.default_action


class PacketFilter(aclgenerator.ACLGenerator):
    """Generates filters and terms from provided policy object."""

    _DEF_MAX_LENGTH = 31
    _PLATFORM = 'packetfilter'
    _DEFAULT_PROTOCOL = 'all'
    SUFFIX = '.pf'
    _TERM = Term

    def _BuildTokens(self) -> tuple[set[str], dict[str, set[str]]]:
        """Build supported tokens for platform.

        Returns:
          tuple containing both supported tokens and sub tokens
        """
        supported_tokens, supported_sub_tokens = super()._BuildTokens()

        supported_tokens |= {'logging', 'destination_interface', 'source_interface'}
        supported_sub_tokens.update(
            {
                'action': {'accept', 'deny', 'reject', 'next'},
                'option': {
                    'established',
                    'tcp-established',
                    'syn',
                    'ack',
                    'fin',
                    'rst',
                    'urg',
                    'psh',
                    'all',
                },
            }
        )

        return supported_tokens, supported_sub_tokens

    def _TranslatePolicy(self, pol: Policy, exp_info: int) -> None:
        self.pf_policies = []
        self.address_book = {}
        self.def_short_to_long = {}
        good_afs = ['inet', 'inet6', 'mixed']
        good_options = ['in', 'out', 'nostate']
        all_protocols_stateful = True

        for header, terms in pol.filters:
            filter_type = None
            filter_options = header.FilterOptions(self._PLATFORM)[1:]
            filter_name = header.FilterName(self._PLATFORM)
            direction = ''

            # ensure all options after the filter name are expected
            for opt in filter_options:
                if opt not in good_afs + good_options:
                    raise aclgenerator.UnsupportedTargetOptionError(
                        '%s %s %s %s'
                        % (
                            '\nUnsupported option found in',
                            self._PLATFORM,
                            'target definition:',
                            opt,
                        )
                    )

            # pf will automatically add 'keep state flags S/SA' to all TCP connections
            # by default.
            if 'nostate' in filter_options:
                all_protocols_stateful = False

            if 'in' in filter_options:
                direction = 'in'
            elif 'out' in filter_options:
                direction = 'out'

            # Check for matching af
            for address_family in good_afs:
                if address_family in filter_options:
                    # should not specify more than one AF in options
                    if filter_type is not None:
                        raise aclgenerator.UnsupportedFilterError(
                            '%s %s %s %s'
                            % (
                                '\nMay only specify one of',
                                good_afs,
                                'in filter options:',
                                filter_options,
                            )
                        )
                    filter_type = address_family
            if filter_type is None:
                filter_type = 'inet'

            # add the terms
            new_terms = []
            term_names = set()

            for term in terms:
                term.name = self.FixTermLength(term.name)
                if term.name in term_names:
                    raise DuplicateTermError(f'You have a duplicate term: {term.name}')
                term_names.add(term.name)

                for source_addr in term.source_address:
                    src_token = source_addr.parent_token[: self._DEF_MAX_LENGTH]

                    if (
                        src_token in self.def_short_to_long
                        and self.def_short_to_long[src_token] != source_addr.parent_token
                    ):
                        raise DuplicateShortenedTableNameError(
                            'There is a shortened name conflict between names %s and %s '
                            '(different named objects would conflict when shortened to %s)'
                            % (
                                self.def_short_to_long[src_token],
                                source_addr.parent_token,
                                src_token,
                            )
                        )
                    else:
                        self.def_short_to_long[src_token] = source_addr.parent_token

                    if src_token not in self.address_book:
                        self.address_book[src_token] = {source_addr}
                    else:
                        self.address_book[src_token].add(source_addr)

                for dest_addr in term.destination_address:
                    dst_token = dest_addr.parent_token[: self._DEF_MAX_LENGTH]

                    if (
                        dst_token in self.def_short_to_long
                        and self.def_short_to_long[dst_token] != dest_addr.parent_token
                    ):
                        raise DuplicateShortenedTableNameError(
                            'There is a shortened name conflict between names %s and %s '
                            '(different named objects would conflict when shortened to %s)'
                            % (
                                self.def_short_to_long[dst_token],
                                dest_addr.parent_token,
                                dst_token,
                            )
                        )
                    else:
                        self.def_short_to_long[dst_token] = dest_addr.parent_token

                    if dst_token not in self.address_book:
                        self.address_book[dst_token] = {dest_addr}
                    else:
                        self.address_book[dst_token].add(dest_addr)

                if not term:
                    continue

                if term.destination_interface and term.source_interface:
                    raise Error(
                        "packetfilter only supports destination interface XOR source interface per term"
                    )

                term_direction = direction
                if term.destination_interface:
                    term_direction = 'out'
                    term.interface = term.destination_interface
                elif term.source_interface:
                    term_direction = 'in'
                    term.interface = term.source_interface
                else:
                    term.interface = None

                new_terms.append(
                    self._TERM(
                        term, filter_name, all_protocols_stateful, filter_type, term_direction
                    )
                )

            self.pf_policies.append((header, filter_name, filter_type, new_terms))

    def __str__(self) -> str:
        """Render the output of the PF policy into config."""
        target = []
        pretty_platform = f'{self._PLATFORM[0].upper()}{self._PLATFORM[1:]}'
        # Create address table.
        for name in sorted(self.address_book):
            v4 = sorted([x for x in self.address_book[name] if x.version == 4])
            v6 = sorted([x for x in self.address_book[name] if x.version == 6])
            entries = ',\\\n'.join(str(x) for x in v4 + v6)
            target.append(f'table <{name}> {{{entries}}}')
        # pylint: disable=unused-variable
        for header, filter_name, filter_type, terms in self.pf_policies:
            # Add comments for this filter
            target.append(f'# {pretty_platform} {header.FilterName(self._PLATFORM)} Policy')

            # reformat long text comments, if needed
            comments = aclgenerator.WrapWords(header.comment, 70)
            if comments and comments[0]:
                for line in comments:
                    target.append(f'# {line}')
                target.append('#')
            # add the p4 tags
            target.extend(aclgenerator.AddRepositoryTags('# '))
            target.append(f"# {filter_type}")

            # add the terms
            for term in terms:
                term_str = str(term)
                if term_str:
                    target.append(term_str)
            target.append('')

        return '\n'.join(target)
