import windows
import ctypes
import socket
import struct

from windows import winproxy
import windows.generated_def as gdef
from windows.com import interfaces as cominterfaces
from windows.generated_def.winstructs import *
from windows.generated_def.windef import *


class TCP4Connection(MIB_TCPROW_OWNER_PID):
    """A TCP4 socket (connected or listening)"""
    @property
    def established(self):
        """``True`` if connection is established else it's a listening socket"""
        return self.dwState == MIB_TCP_STATE_ESTAB

    @property
    def remote_port(self):
        """:type: :class:`int`"""
        if not self.established:
            return None
        return socket.ntohs(self.dwRemotePort)

    @property
    def local_port(self):
        """:type: :class:`int`"""
        return socket.ntohs(self.dwLocalPort)

    @property
    def local_addr(self):
        """Local address IP (x.x.x.x)

        :type: :class:`str`"""
        return socket.inet_ntoa(struct.pack("<I", self.dwLocalAddr))

    @property
    def remote_addr(self):
        """remote address IP (x.x.x.x)

        :type: :class:`str`"""
        if not self.established:
            return None
        return socket.inet_ntoa(struct.pack("<I", self.dwRemoteAddr))

    @property
    def remote_proto(self):
        """Identification of the protocol associated with the remote port.
           Equals ``remote_port`` if no protocol is associated with it.

           :type: :class:`str` or :class:`int`
        """
        try:
            return socket.getservbyport(self.remote_port, 'tcp')
        except socket.error:
            return self.remote_port

    @property
    def remote_host(self):
        """Identification of the remote hostname.
           Equals ``remote_addr`` if the resolution fails

           :type: :class:`str` or :class:`int`
        """

        try:
            return socket.gethostbyaddr(self.remote_addr)
        except socket.error:
            return self.remote_addr

    def close(self):
        """Close the connection <require elevated process>"""
        closing = MIB_TCPROW()
        closing.dwState = MIB_TCP_STATE_DELETE_TCB
        closing.dwLocalAddr = self.dwLocalAddr
        closing.dwLocalPort = self.dwLocalPort
        closing.dwRemoteAddr = self.dwRemoteAddr
        closing.dwRemotePort = self.dwRemotePort
        return winproxy.SetTcpEntry(ctypes.byref(closing))

    def __repr__(self):
        if not self.established:
            return "<TCP IPV4 Listening socket on {0}:{1}>".format(self.local_addr, self.local_port)
        return "<TCP IPV4 Connection {s.local_addr}:{s.local_port} -> {s.remote_addr}:{s.remote_port}>".format(s=self)


class TCP6Connection(MIB_TCP6ROW_OWNER_PID):
    """A TCP6 socket (connected or listening)"""
    @staticmethod
    def _str_ipv6_addr(addr):
        return ":".join(c.encode('hex') for c in addr)

    @property
    def established(self):
        """``True`` if connection is established else it's a listening socket"""
        return self.dwState == MIB_TCP_STATE_ESTAB

    @property
    def remote_port(self):
        """:type: :class:`int`"""
        if not self.established:
            return None
        return socket.ntohs(self.dwRemotePort)

    @property
    def local_port(self):
        """:type: :class:`int`"""
        return socket.ntohs(self.dwLocalPort)

    @property
    def local_addr(self):
        """Local address IP

        :type: :class:`str`"""
        return self._str_ipv6_addr(self.ucLocalAddr)

    @property
    def remote_addr(self):
        """remote address IP

        :type: :class:`str`"""
        if not self.established:
            return None
        return self._str_ipv6_addr(self.ucRemoteAddr)

    @property
    def remote_proto(self):
        """Equals to ``self.remote_port`` for Ipv6"""
        return self.remote_port

    @property
    def remote_host(self):
        """Equals to ``self.remote_addr`` for Ipv6"""
        return self.remote_addr

    def close(self):
        raise NotImplementedError("Closing IPV6 connection non implemented")

    def __repr__(self):
        if not self.established:
            return "<TCP IPV6 Listening socket on {0}:{1}>".format(self.local_addr, self.local_port)
        return "<TCP IPV6 Connection {0}:{1} -> {2}:{3}>".format(self.local_addr, self.local_port, self.remote_addr, self.remote_port)


def get_MIB_TCPTABLE_OWNER_PID_from_buffer(buffer):
    x = windows.generated_def.winstructs.MIB_TCPTABLE_OWNER_PID.from_buffer(buffer)
    nb_entry = x.dwNumEntries

    class _GENERATED_MIB_TCPTABLE_OWNER_PID(ctypes.Structure):
            _fields_ = [
                ("dwNumEntries", DWORD),
                ("table", TCP4Connection * nb_entry),
            ]

    return _GENERATED_MIB_TCPTABLE_OWNER_PID.from_buffer(buffer)


def get_MIB_TCP6TABLE_OWNER_PID_from_buffer(buffer):
    x = windows.generated_def.winstructs.MIB_TCP6TABLE_OWNER_PID.from_buffer(buffer)
    nb_entry = x.dwNumEntries

    # Struct _MIB_TCP6TABLE_OWNER_PID definitions
    class _GENERATED_MIB_TCP6TABLE_OWNER_PID(Structure):
            _fields_ = [
                ("dwNumEntries", DWORD),
                ("table", TCP6Connection * nb_entry),
            ]

    return _GENERATED_MIB_TCP6TABLE_OWNER_PID.from_buffer(buffer)

class Firewall(cominterfaces.INetFwPolicy2):
    """The windows firewall"""
    @property
    def rules(self):
        """The rules of the firewall

        :type: [:class:`FirewallRule`] -- A list of rule
        """
        ifw_rules = cominterfaces.INetFwRules()
        self.get_Rules(ifw_rules)

        nb_rules = gdef.LONG()
        ifw_rules.get_Count(nb_rules)

        unknw = cominterfaces.IUnknown()
        ifw_rules.get__NewEnum(unknw)

        pVariant = cominterfaces.IEnumVARIANT()
        unknw.QueryInterface(pVariant.IID, pVariant)

        count = gdef.ULONG()
        var = windows.com.ImprovedVariant()

        rules = []
        for i in range(nb_rules.value):
            pVariant.Next(1, var, count)
            if not count.value:
                break
            rule = FirewallRule()
            idisp = var.asdispatch
            idisp.QueryInterface(rule.IID, rule)
            rules.append(rule)
        return rules

    @property
    def current_profile_types(self):
        """Mask of the profiles currently enabled

        :type: :class:`long`
        """
        cpt = gdef.LONG()
        self.get_CurrentProfileTypes(cpt)
        return cpt.value

    @property
    def enabled(self):
        """A maping of the active firewall profiles

        {

        ``NET_FW_PROFILE_TYPE2_.NET_FW_PROFILE2_DOMAIN(0x1L)``: ``True`` or ``False``,

        ``NET_FW_PROFILE_TYPE2_.NET_FW_PROFILE2_PRIVATE(0x2L)``: ``True`` or ``False``,

        ``NET_FW_PROFILE_TYPE2_.NET_FW_PROFILE2_PUBLIC(0x4L)``: ``True`` or ``False``,

        }


        :type: :class:`dict`
        """
        profiles = [gdef.NET_FW_PROFILE2_DOMAIN, gdef.NET_FW_PROFILE2_PRIVATE, gdef.NET_FW_PROFILE2_PUBLIC]
        return {prof: self.enabled_for_profile_type(prof) for prof in profiles}


    def enabled_for_profile_type(self, profile_type):
        enabled = gdef.VARIANT_BOOL()
        self.get_FirewallEnabled(profile_type, enabled)
        return enabled.value



class FirewallRule(cominterfaces.INetFwRule):
    """A rule of the firewall"""
    @property
    def name(self):
        """Name of the rule

        :type: :class:`unicode`
        """
        name = gdef.BSTR()
        self.get_Name(name)
        return name.value

    @property
    def description(self):
        """Description of the rule

        :type: :class:`unicode`
        """
        description = gdef.BSTR()
        self.get_Description(description)
        return description.value

    @property
    def application_name(self):
        """Name of the application to which apply the rule

        :type: :class:`unicode`
        """
        applicationname = gdef.BSTR()
        self.get_ApplicationName(applicationname)
        return applicationname.value

    @property
    def service_name(self):
        """Name of the service to which apply the rule

        :type: :class:`unicode`
        """
        servicename = gdef.BSTR()
        self.get_ServiceName(servicename)
        return servicename.value

    @property
    def protocol(self):
        """Protocol to which apply the rule

        :type: :class:`long`
        """
        protocol = gdef.LONG()
        self.get_Protocol(protocol)
        return protocol.value

    @property
    def local_address(self):
        """Local address of the rule

        :type: :class:`unicode`
        """
        local_address = gdef.BSTR()
        self.get_LocalAddresses(local_address)
        return local_address.value

    @property
    def remote_address(self):
        """Remote address of the rule

        :type: :class:`unicode`
        """
        remote_address = gdef.BSTR()
        self.get_RemoteAddresses(remote_address)
        return remote_address.value

    @property
    def direction(self):
        """Direction of the rule, values might be:

            * ``NET_FW_RULE_DIRECTION_.NET_FW_RULE_DIR_IN(0x1L)``
            * ``NET_FW_RULE_DIRECTION_.NET_FW_RULE_DIR_OUT(0x2L)``

        subclass of :class:`long`
        """
        direction = gdef.NET_FW_RULE_DIRECTION()
        self.get_Direction(direction)
        return direction.value

    @property
    def interface_types(self):
        """Types of interface of the rule

        :type: :class:`unicode`
        """
        interface_type = gdef.BSTR()
        self.get_InterfaceTypes(interface_type)
        return interface_type.value

    @property
    def local_port(self):
        """Local port of the rule

        :type: :class:`unicode`
        """
        local_port = gdef.BSTR()
        self.get_LocalPorts(local_port)
        return local_port.value

    @property
    def remote_port(self):
        """Remote port of the rule

        :type: :class:`unicode`
        """
        remote_port = gdef.BSTR()
        self.get_RemotePorts(remote_port)
        return remote_port.value

    @property
    def action(self):
        """Action of the rule, values might be:

            * ``NET_FW_ACTION_.NET_FW_ACTION_BLOCK(0x0L)``
            * ``NET_FW_ACTION_.NET_FW_ACTION_ALLOW(0x1L)``

        subclass of :class:`long`
        """
        action = gdef.NET_FW_ACTION()
        self.get_Action(action)
        return action.value

    @property
    def enabled(self):
        """``True`` if rule is enabled"""
        enabled = gdef.VARIANT_BOOL()
        self.get_Enabled(enabled)
        return enabled.value

    @property
    def grouping(self):
        """Grouping of the rule

        :type: :class:`unicode`
        """
        grouping = gdef.BSTR()
        self.get_RemotePorts(grouping)
        return grouping.value

    @property
    def icmp_type_and_code(self):
        icmp_type_and_code = gdef.BSTR()
        self.get_RemotePorts(icmp_type_and_code)
        return icmp_type_and_code.value

    def __repr__(self):
        return u'<{0} "{1}">'.format(type(self).__name__, self.name).encode("ascii", errors='backslashreplace')

class Network(object):
    NetFwPolicy2 = windows.com.IID.from_string("E2B3C97F-6AE1-41AC-817A-F6F92166D7DD")

    @property
    def firewall(self):
        """The firewall of the system

        :type: :class:`Firewall`
        """
        windows.com.init()
        firewall = Firewall()
        windows.com.create_instance(self.NetFwPolicy2, firewall)
        return firewall

    @staticmethod
    def _get_tcp_ipv4_sockets():
        size = ctypes.c_uint(0)
        try:
            winproxy.GetExtendedTcpTable(None, ctypes.byref(size), ulAf=AF_INET)
        except winproxy.IphlpapiError:
            pass  # Allow us to set size to the needed value
        buffer = (ctypes.c_char * size.value)()
        winproxy.GetExtendedTcpTable(buffer, ctypes.byref(size), ulAf=AF_INET)
        t = get_MIB_TCPTABLE_OWNER_PID_from_buffer(buffer)
        return list(t.table)

    @staticmethod
    def _get_tcp_ipv6_sockets():
        size = ctypes.c_uint(0)
        try:
            winproxy.GetExtendedTcpTable(None, ctypes.byref(size), ulAf=AF_INET6)
        except winproxy.IphlpapiError:
            pass  # Allow us to set size to the needed value
        buffer = (ctypes.c_char * size.value)()
        winproxy.GetExtendedTcpTable(buffer, ctypes.byref(size), ulAf=AF_INET6)
        t = get_MIB_TCP6TABLE_OWNER_PID_from_buffer(buffer)
        return list(t.table)


    ipv4 = property(lambda self: self._get_tcp_ipv4_sockets())
    """List of TCP IPv4 socket (connection and listening)

        :type: [:class:`TCP4Connection`]"""

    ipv6 = property(lambda self: self._get_tcp_ipv6_sockets())
    """List of TCP IPv6 socket (connection and listening)

      :type: [:class:`TCP6Connection`]
    """

[((4444, 4519), 'windows.generated_def.winstructs.MIB_TCPTABLE_OWNER_PID.from_buffer', 'windows.generated_def.winstructs.MIB_TCPTABLE_OWNER_PID.from_buffer', (['buffer'], {}), '(buffer)\n', (4511, 4519), False, 'import windows\n'), ((4877, 4953), 'windows.generated_def.winstructs.MIB_TCP6TABLE_OWNER_PID.from_buffer', 'windows.generated_def.winstructs.MIB_TCP6TABLE_OWNER_PID.from_buffer', (['buffer'], {}), '(buffer)\n', (4945, 4953), False, 'import windows\n'), ((11219, 11286), 'windows.com.IID.from_string', 'windows.com.IID.from_string', (['"""E2B3C97F-6AE1-41AC-817A-F6F92166D7DD"""'], {}), "('E2B3C97F-6AE1-41AC-817A-F6F92166D7DD')\n", (11246, 11286), False, 'import windows\n'), ((682, 713), 'socket.ntohs', 'socket.ntohs', (['self.dwRemotePort'], {}), '(self.dwRemotePort)\n', (694, 713), False, 'import socket\n'), ((804, 834), 'socket.ntohs', 'socket.ntohs', (['self.dwLocalPort'], {}), '(self.dwLocalPort)\n', (816, 834), False, 'import socket\n'), ((3211, 3242), 'socket.ntohs', 'socket.ntohs', (['self.dwRemotePort'], {}), '(self.dwRemotePort)\n', (3223, 3242), False, 'import socket\n'), ((3333, 3363), 'socket.ntohs', 'socket.ntohs', (['self.dwLocalPort'], {}), '(self.dwLocalPort)\n', (3345, 3363), False, 'import socket\n'), ((5532, 5559), 'windows.com.interfaces.INetFwRules', 'cominterfaces.INetFwRules', ([], {}), '()\n', (5557, 5559), True, 'from windows.com import interfaces as cominterfaces\n'), ((5614, 5625), 'windows.generated_def.LONG', 'gdef.LONG', ([], {}), '()\n', (5623, 5625), True, 'import windows.generated_def as gdef\n'), ((5681, 5705), 'windows.com.interfaces.IUnknown', 'cominterfaces.IUnknown', ([], {}), '()\n', (5703, 5705), True, 'from windows.com import interfaces as cominterfaces\n'), ((5764, 5792), 'windows.com.interfaces.IEnumVARIANT', 'cominterfaces.IEnumVARIANT', ([], {}), '()\n', (5790, 5792), True, 'from windows.com import interfaces as cominterfaces\n'), ((5863, 5875), 'windows.generated_def.ULONG', 'gdef.ULONG', ([], {}), '()\n', (5873, 5875), True, 'import windows.generated_def as gdef\n'), ((5890, 5919), 'windows.com.ImprovedVariant', 'windows.com.ImprovedVariant', ([], {}), '()\n', (5917, 5919), False, 'import windows\n'), ((6403, 6414), 'windows.generated_def.LONG', 'gdef.LONG', ([], {}), '()\n', (6412, 6414), True, 'import windows.generated_def as gdef\n'), ((7164, 7183), 'windows.generated_def.VARIANT_BOOL', 'gdef.VARIANT_BOOL', ([], {}), '()\n', (7181, 7183), True, 'import windows.generated_def as gdef\n'), ((7473, 7484), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (7482, 7484), True, 'import windows.generated_def as gdef\n'), ((7683, 7694), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (7692, 7694), True, 'import windows.generated_def as gdef\n'), ((7947, 7958), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (7956, 7958), True, 'import windows.generated_def as gdef\n'), ((8211, 8222), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (8220, 8222), True, 'import windows.generated_def as gdef\n'), ((8442, 8453), 'windows.generated_def.LONG', 'gdef.LONG', ([], {}), '()\n', (8451, 8453), True, 'import windows.generated_def as gdef\n'), ((8670, 8681), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (8679, 8681), True, 'import windows.generated_def as gdef\n'), ((8917, 8928), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (8926, 8928), True, 'import windows.generated_def as gdef\n'), ((9306, 9334), 'windows.generated_def.NET_FW_RULE_DIRECTION', 'gdef.NET_FW_RULE_DIRECTION', ([], {}), '()\n', (9332, 9334), True, 'import windows.generated_def as gdef\n'), ((9562, 9573), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (9571, 9573), True, 'import windows.generated_def as gdef\n'), ((9799, 9810), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (9808, 9810), True, 'import windows.generated_def as gdef\n'), ((10027, 10038), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (10036, 10038), True, 'import windows.generated_def as gdef\n'), ((10382, 10402), 'windows.generated_def.NET_FW_ACTION', 'gdef.NET_FW_ACTION', ([], {}), '()\n', (10400, 10402), True, 'import windows.generated_def as gdef\n'), ((10561, 10580), 'windows.generated_def.VARIANT_BOOL', 'gdef.VARIANT_BOOL', ([], {}), '()\n', (10578, 10580), True, 'import windows.generated_def as gdef\n'), ((10779, 10790), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (10788, 10790), True, 'import windows.generated_def as gdef\n'), ((10938, 10949), 'windows.generated_def.BSTR', 'gdef.BSTR', ([], {}), '()\n', (10947, 10949), True, 'import windows.generated_def as gdef\n'), ((11418, 11436), 'windows.com.init', 'windows.com.init', ([], {}), '()\n', (11434, 11436), False, 'import windows\n'), ((11475, 11531), 'windows.com.create_instance', 'windows.com.create_instance', (['self.NetFwPolicy2', 'firewall'], {}), '(self.NetFwPolicy2, firewall)\n', (11502, 11531), False, 'import windows\n'), ((11623, 11639), 'ctypes.c_uint', 'ctypes.c_uint', (['(0)'], {}), '(0)\n', (11636, 11639), False, 'import ctypes\n'), ((12116, 12132), 'ctypes.c_uint', 'ctypes.c_uint', (['(0)'], {}), '(0)\n', (12129, 12132), False, 'import ctypes\n'), ((978, 1013), 'struct.pack', 'struct.pack', (['"""<I"""', 'self.dwLocalAddr'], {}), "('<I', self.dwLocalAddr)\n", (989, 1013), False, 'import struct\n'), ((1217, 1253), 'struct.pack', 'struct.pack', (['"""<I"""', 'self.dwRemoteAddr'], {}), "('<I', self.dwRemoteAddr)\n", (1228, 1253), False, 'import struct\n'), ((1537, 1582), 'socket.getservbyport', 'socket.getservbyport', (['self.remote_port', '"""tcp"""'], {}), "(self.remote_port, 'tcp')\n", (1557, 1582), False, 'import socket\n'), ((1891, 1929), 'socket.gethostbyaddr', 'socket.gethostbyaddr', (['self.remote_addr'], {}), '(self.remote_addr)\n', (1911, 1929), False, 'import socket\n'), ((2389, 2410), 'ctypes.byref', 'ctypes.byref', (['closing'], {}), '(closing)\n', (2401, 2410), False, 'import ctypes\n'), ((11927, 11945), 'ctypes.byref', 'ctypes.byref', (['size'], {}), '(size)\n', (11939, 11945), False, 'import ctypes\n'), ((12421, 12439), 'ctypes.byref', 'ctypes.byref', (['size'], {}), '(size)\n', (12433, 12439), False, 'import ctypes\n'), ((11700, 11718), 'ctypes.byref', 'ctypes.byref', (['size'], {}), '(size)\n', (11712, 11718), False, 'import ctypes\n'), ((12193, 12211), 'ctypes.byref', 'ctypes.byref', (['size'], {}), '(size)\n', (12205, 12211), False, 'import ctypes\n')]