# This file is part of Volatility 3 and is released under the Volatility Software License 1.0
# Available at https://www.volatilityfoundation.org/license/vsl-v1.0

import logging
from typing import Iterable, Callable, Tuple

from volatility3.framework import exceptions, renderers, interfaces, constants
from volatility3.framework.configuration import requirements
from volatility3.framework.interfaces import plugins
from volatility3.framework.objects import utility
from volatility3.framework.renderers import format_hints
from volatility3.framework.symbols import linux
from volatility3.plugins.linux import pslist

vollog = logging.getLogger(__name__)


class Netstat(plugins.PluginInterface):
    """Lists all network connections for all processes."""

    _required_framework_version = (2, 0, 0)

    @classmethod
    def get_requirements(cls):
        return [
            requirements.ModuleRequirement(
                name="kernel",
                description="Linux kernel module",
                architectures=["Intel32", "Intel64", "AArch64", "AMD64"],
            ),
            requirements.PluginRequirement(
                name="pslist", plugin=pslist.PsList, version=(2, 0, 0)
            ),
            requirements.VersionRequirement(
                name="linuxutils", component=linux.LinuxUtilities, version=(2, 1, 0)
            ),
            requirements.ListRequirement(
                name="pid",
                description="Filter on specific process IDs",
                element_type=int,
                optional=True,
            ),
        ]

    @classmethod
    def list_fds_no_path(
        cls,
        context: interfaces.context.ContextInterface,
        vmlinux_name: str,
        filter_func: Callable[[int], bool] = lambda _: False,
    ) -> Iterable[Tuple[int, str, interfaces.objects.ObjectInterface, Tuple[int, interfaces.objects.ObjectInterface]]]:
        """
        Lists file descriptors for all processes without retrieving the file path.

        Yields:
            A tuple containing:
                - PID
                - Process command name
                - Task object
                - Tuple of (fd number, file object)
        """
        for task in pslist.PsList.list_tasks(context, vmlinux_name, filter_func):
            pid = int(task.pid)
            if filter_func(pid):
                continue

            task_comm = utility.array_to_string(task.comm)

            files = task.files
            if not files:
                continue

            try:
                fd_table = files.get_fds()
            except exceptions.InvalidAddressException:
                continue

            if not fd_table:
                continue

            for fd_num, filp in fd_table.items():
                if not filp:
                    continue

                yield pid, task_comm, task, (fd_num, filp)

    @classmethod
    def list_sockets(
        cls,
        context: interfaces.context.ContextInterface,
        kernel_module_name: str,
        filter_func: Callable[[int], bool] = lambda _: False,
    ) -> Iterable[Tuple[str, int, interfaces.objects.ObjectInterface]]:
        vmlinux = context.modules[kernel_module_name]

        # Try to get the socket_file_ops symbol
        try:
            socket_file_ops = vmlinux.object_from_symbol("socket_file_ops")
        except exceptions.SymbolError:
            vollog.warning("Symbol 'socket_file_ops' not found. The plugin may not work correctly.")
            socket_file_ops = None

        # Use the custom list_fds_no_path method
        fd_generator = cls.list_fds_no_path(context, vmlinux.name, filter_func)

        for pid, task_comm, task, fd_info in fd_generator:
            fd_num, filp = fd_info

            # Check if filp.f_op points to socket_file_ops
            try:
                if socket_file_ops and filp.f_op == socket_file_ops:
                    # It's a socket
                    pass
                else:
                    continue
            except exceptions.InvalidAddressException:
                continue

            # Get the socket object from filp
            try:
                socket = filp.private_data.dereference().cast("socket")
            except exceptions.InvalidAddressException:
                continue

            # Ensure the socket object is valid
            if not context.layers[task.vol.native_layer_name].is_valid(
                socket.vol.offset, socket.vol.size
            ):
                continue

            task_name = utility.array_to_string(task.comm)
            yield task_name, pid, socket

    def _generator(self):
        filter_func = pslist.PsList.create_pid_filter(self.config.get("pid", None))

        for task_name, pid, socket in self.list_sockets(
            self.context, self.config["kernel"], filter_func=filter_func
        ):
            # Obtain the address family
            try:
                family = socket.ops.family
            except exceptions.InvalidAddressException:
                continue

            # Handle UNIX sockets (AF_UNIX)
            if family == socket.vol.consts.AF_UNIX:
                try:
                    unix_sock = socket.sk.dereference().cast("unix_sock")
                    if not unix_sock:
                        continue
                    # Get the path from the unix_sock structure
                    addr = unix_sock.addr
                    if addr:
                        path = utility.array_to_string(addr.name, encoding="utf-8")
                    else:
                        path = ""
                except exceptions.InvalidAddressException:
                    continue

                yield (
                    0,
                    (
                        format_hints.Hex(socket.vol.offset),
                        "UNIX",
                        path,
                        0,
                        "",
                        0,
                        "",
                        f"{task_name}/{pid}",
                    ),
                )

            # Handle IPv4 and IPv6 sockets
            elif family in [socket.vol.consts.AF_INET, socket.vol.consts.AF_INET6]:
                try:
                    # Get protocol (TCP/UDP)
                    protocol = socket.sk.__sk_common.skc_protocol
                    if protocol == socket.vol.consts.IPPROTO_TCP:
                        proto = "TCP"
                    elif protocol == socket.vol.consts.IPPROTO_UDP:
                        proto = "UDP"
                    else:
                        proto = f"IPPROTO_{protocol}"

                    # Get state
                    state = socket.sk.__sk_common.skc_state

                    # Get IP addresses and ports
                    if family == socket.vol.consts.AF_INET:
                        lip = utility.array_to_ip(socket.sk.__sk_common.skc_rcv_saddr, family="ipv4")
                        rip = utility.array_to_ip(socket.sk.__sk_common.skc_daddr, family="ipv4")
                    else:
                        lip = utility.array_to_ip(socket.sk.__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr8, family="ipv6")
                        rip = utility.array_to_ip(socket.sk.__sk_common.skc_v6_daddr.in6_u.u6_addr8, family="ipv6")

                    lport = socket.sk.__sk_common.skc_num
                    rport = socket.sk.__sk_common.skc_dport

                    # Convert network byte order to host byte order for ports
                    rport = int.from_bytes(rport.to_bytes(2, byteorder='little'), byteorder='big')

                except (exceptions.InvalidAddressException, AttributeError, ValueError):
                    continue

                yield (
                    0,
                    (
                        format_hints.Hex(socket.vol.offset),
                        proto,
                        lip,
                        lport,
                        rip,
                        rport,
                        str(state),
                        f"{task_name}/{pid}",
                    ),
                )

    def run(self):
        return renderers.TreeGrid(
            [
                ("Offset", format_hints.Hex),
                ("Proto", str),
                ("Local IP", str),
                ("Local Port", int),
                ("Remote IP", str),
                ("Remote Port", int),
                ("State", str),
                ("Process", str),
            ],
            self._generator(),
        )
