from __future__ import annotations

import textwrap
from io import BytesIO
from typing import TYPE_CHECKING

from flow.record.fieldtypes import datetime

from dissect.target.plugins.os.unix.linux.iptables import IptablesSavePlugin

if TYPE_CHECKING:
    from dissect.target.filesystem import VirtualFilesystem
    from dissect.target.target import Target


def test_iptables_plugin(target_unix_users: Target, fs_unix: VirtualFilesystem) -> None:
    ipv4_rules_path = "/etc/iptables/rules.v4"
    ipv4_rules = """# Generated by iptables-save v1.8.7 on Wed May 31 13:37:00 2023
    *filter
    :INPUT ACCEPT [0:0]
    :FORWARD DROP [0:0]
    :OUTPUT ACCEPT [0:0]
    :DOCKER - [0:0]
    :DOCKER-ISOLATION-STAGE-1 - [0:0]
    :DOCKER-ISOLATION-STAGE-2 - [0:0]
    :DOCKER-USER - [0:0]
    -A FORWARD -j DOCKER-USER
    -A DOCKER-ISOLATION-STAGE-1 -i br-aaaaaaaaaaaa ! -o br-bbbbbbbbbbbb -j DOCKER-ISOLATION-STAGE-2
    COMMIT
    # Completed on Wed May 31 13:37:00 2023
    # Generated by iptables-save v1.8.7 on Wed May 31 13:37:00 2023
    *nat
    :PREROUTING ACCEPT [0:0]
    :INPUT ACCEPT [0:0]
    :OUTPUT ACCEPT [0:0]
    :POSTROUTING ACCEPT [0:0]
    :DOCKER - [0:0]
    -A PREROUTING -m addrtype --dst-type LOCAL -j DOCKER
    -A DOCKER -d 127.0.0.1/32 ! -i br-aaaaaaaaaaaa -p tcp -m tcp --dport 3456 -j DNAT --to-destination 172.18.0.2:3456
    COMMIT
    # Completed on Wed May 31 13:37:00 2023
    """

    ipv6_rules_path = "/etc/iptables/rules.v6"
    ipv6_rules = """# Generated by ip6tables-save v1.8.7 on Wed Jun 1 13:37:00 2023
    *filter
    :INPUT ACCEPT [21:2975]
    :FORWARD ACCEPT [0:0]
    :OUTPUT ACCEPT [0:0]
    [78:2323] -A PREROUTING -m addrtype --dst-type LOCAL -j DOCKER
    COMMIT
    # Completed on Wed Jun 1 13:37:00 2023
    """

    fs_unix.map_file_fh(
        ipv4_rules_path,
        BytesIO(textwrap.dedent(ipv4_rules).encode()),
    )

    fs_unix.map_file_fh(
        ipv6_rules_path,
        BytesIO(textwrap.dedent(ipv6_rules).encode()),
    )

    target_unix_users.add_plugin(IptablesSavePlugin)
    results = list(target_unix_users.iptables())

    assert len(results) == 20

    assert results[0].ts == datetime(2023, 5, 31, 13, 37, 0)
    assert results[0].program == "iptables-save"
    assert results[0].version == "1.8.7"
    assert results[0].table == "filter"
    assert results[0].chain == "INPUT"
    assert results[0].type == "policy"
    assert results[0].rule == "ACCEPT"
    assert results[0].packet_count == 0
    assert results[0].byte_count == 0

    assert results[5].ts == datetime(2023, 5, 31, 13, 37, 0)
    assert results[5].program == "iptables-save"
    assert results[5].version == "1.8.7"
    assert results[5].table == "filter"
    assert results[5].chain == "DOCKER-ISOLATION-STAGE-2"
    assert results[5].type == "policy"
    assert results[5].rule == "-"
    assert results[5].packet_count == 0
    assert results[5].byte_count == 0

    assert results[15].ts == datetime(2023, 5, 31, 13, 37, 0)
    assert results[15].program == "iptables-save"
    assert results[15].version == "1.8.7"
    assert results[15].table == "nat"
    assert results[15].chain == "DOCKER"
    assert results[15].type == "rule"
    assert results[15].rule == (
        "-A DOCKER -d 127.0.0.1/32 ! -i br-aaaaaaaaaaaa -p tcp -m tcp"
        " --dport 3456 -j DNAT --to-destination 172.18.0.2:3456"
    )
    assert results[15].packet_count is None
    assert results[15].byte_count is None

    assert results[16].ts == datetime(2023, 6, 1, 13, 37, 0)
    assert results[16].program == "ip6tables-save"
    assert results[16].version == "1.8.7"
    assert results[16].table == "filter"
    assert results[16].chain == "INPUT"
    assert results[16].type == "policy"
    assert results[16].rule == "ACCEPT"
    assert results[16].packet_count == 21
    assert results[16].byte_count == 2975

    assert results[19].ts == datetime(2023, 6, 1, 13, 37, 0)
    assert results[19].program == "ip6tables-save"
    assert results[19].version == "1.8.7"
    assert results[19].table == "filter"
    assert results[19].chain == "PREROUTING"
    assert results[19].type == "rule"
    assert results[19].rule == "-A PREROUTING -m addrtype --dst-type LOCAL -j DOCKER"
    assert results[19].packet_count == 78
    assert results[19].byte_count == 2323


def test_iptables_plugin_ufw(target_unix_users: Target, fs_unix: VirtualFilesystem) -> None:
    ufw_rules_path = "/etc/ufw/user.rules"

    ufw_rules = """*filter
    :ufw-user-input - [0:0]
    :ufw-user-output - [0:0]
    :ufw-user-forward - [0:0]
    :ufw-user-limit - [0:0]
    :ufw-user-limit-accept - [0:0]
    ### RULES ###
    -A ufw-user-limit -m limit --limit 3/minute -j LOG --log-prefix "[UFW LIMIT BLOCK] "
    -A ufw-user-limit -j REJECT
    -A ufw-user-limit-accept -j ACCEPT
    COMMIT
    """

    fs_unix.map_file_fh(
        ufw_rules_path,
        BytesIO(textwrap.dedent(ufw_rules).encode()),
    )

    target_unix_users.add_plugin(IptablesSavePlugin)
    results = list(target_unix_users.iptables())

    # 5 policies, 3 rules
    assert len(results) == 5 + 3

    assert results[-1].source == ufw_rules_path
    assert results[-1].rule == "-A ufw-user-limit-accept -j ACCEPT"
    assert results[-1].table == "filter"
    assert results[-1].chain == "ufw-user-limit-accept"
    assert results[-1].type == "rule"
