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`]
    """

[((148, 8, 148, 83), 'windows.generated_def.winstructs.MIB_TCPTABLE_OWNER_PID.from_buffer', 'windows.generated_def.winstructs.MIB_TCPTABLE_OWNER_PID.from_buffer', ({(148, 76, 148, 82): 'buffer'}, {}), '(buffer)', False, 'import windows\n'), ((161, 8, 161, 84), 'windows.generated_def.winstructs.MIB_TCP6TABLE_OWNER_PID.from_buffer', 'windows.generated_def.winstructs.MIB_TCP6TABLE_OWNER_PID.from_buffer', ({(161, 77, 161, 83): 'buffer'}, {}), '(buffer)', False, 'import windows\n'), ((400, 19, 400, 86), 'windows.com.IID.from_string', 'windows.com.IID.from_string', ({(400, 47, 400, 85): '"""E2B3C97F-6AE1-41AC-817A-F6F92166D7DD"""'}, {}), "('E2B3C97F-6AE1-41AC-817A-F6F92166D7DD')", False, 'import windows\n'), ((25, 15, 25, 46), 'socket.ntohs', 'socket.ntohs', ({(25, 28, 25, 45): 'self.dwRemotePort'}, {}), '(self.dwRemotePort)', False, 'import socket\n'), ((30, 15, 30, 45), 'socket.ntohs', 'socket.ntohs', ({(30, 28, 30, 44): 'self.dwLocalPort'}, {}), '(self.dwLocalPort)', False, 'import socket\n'), ((105, 15, 105, 46), 'socket.ntohs', 'socket.ntohs', ({(105, 28, 105, 45): 'self.dwRemotePort'}, {}), '(self.dwRemotePort)', False, 'import socket\n'), ((110, 15, 110, 45), 'socket.ntohs', 'socket.ntohs', ({(110, 28, 110, 44): 'self.dwLocalPort'}, {}), '(self.dwLocalPort)', False, 'import socket\n'), ((181, 20, 181, 47), 'windows.com.interfaces.INetFwRules', 'cominterfaces.INetFwRules', ({}, {}), '()', True, 'from windows.com import interfaces as cominterfaces\n'), ((184, 19, 184, 30), 'windows.generated_def.LONG', 'gdef.LONG', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((187, 16, 187, 40), 'windows.com.interfaces.IUnknown', 'cominterfaces.IUnknown', ({}, {}), '()', True, 'from windows.com import interfaces as cominterfaces\n'), ((190, 19, 190, 47), 'windows.com.interfaces.IEnumVARIANT', 'cominterfaces.IEnumVARIANT', ({}, {}), '()', True, 'from windows.com import interfaces as cominterfaces\n'), ((193, 16, 193, 28), 'windows.generated_def.ULONG', 'gdef.ULONG', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((194, 14, 194, 43), 'windows.com.ImprovedVariant', 'windows.com.ImprovedVariant', ({}, {}), '()', False, 'import windows\n'), ((213, 14, 213, 25), 'windows.generated_def.LONG', 'gdef.LONG', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((239, 18, 239, 37), 'windows.generated_def.VARIANT_BOOL', 'gdef.VARIANT_BOOL', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((253, 15, 253, 26), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((263, 22, 263, 33), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((273, 26, 273, 37), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((283, 22, 283, 33), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((293, 19, 293, 30), 'windows.generated_def.LONG', 'gdef.LONG', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((303, 24, 303, 35), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((313, 25, 313, 36), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((326, 20, 326, 48), 'windows.generated_def.NET_FW_RULE_DIRECTION', 'gdef.NET_FW_RULE_DIRECTION', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((336, 25, 336, 36), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((346, 21, 346, 32), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((356, 22, 356, 33), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((369, 17, 369, 37), 'windows.generated_def.NET_FW_ACTION', 'gdef.NET_FW_ACTION', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((376, 18, 376, 37), 'windows.generated_def.VARIANT_BOOL', 'gdef.VARIANT_BOOL', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((386, 19, 386, 30), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((392, 29, 392, 40), 'windows.generated_def.BSTR', 'gdef.BSTR', ({}, {}), '()', True, 'import windows.generated_def as gdef\n'), ((408, 8, 408, 26), 'windows.com.init', 'windows.com.init', ({}, {}), '()', False, 'import windows\n'), ((410, 8, 410, 64), 'windows.com.create_instance', 'windows.com.create_instance', ({(410, 36, 410, 53): 'self.NetFwPolicy2', (410, 55, 410, 63): 'firewall'}, {}), '(self.NetFwPolicy2, firewall)', False, 'import windows\n'), ((415, 15, 415, 31), 'ctypes.c_uint', 'ctypes.c_uint', ({(415, 29, 415, 30): '0'}, {}), '(0)', False, 'import ctypes\n'), ((427, 15, 427, 31), 'ctypes.c_uint', 'ctypes.c_uint', ({(427, 29, 427, 30): '0'}, {}), '(0)', False, 'import ctypes\n'), ((37, 32, 37, 67), 'struct.pack', 'struct.pack', ({(37, 44, 37, 48): '"""<I"""', (37, 50, 37, 66): 'self.dwLocalAddr'}, {}), "('<I', self.dwLocalAddr)", False, 'import struct\n'), ((46, 32, 46, 68), 'struct.pack', 'struct.pack', ({(46, 44, 46, 48): '"""<I"""', (46, 50, 46, 67): 'self.dwRemoteAddr'}, {}), "('<I', self.dwRemoteAddr)", False, 'import struct\n'), ((56, 19, 56, 64), 'socket.getservbyport', 'socket.getservbyport', ({(56, 40, 56, 56): 'self.remote_port', (56, 58, 56, 63): '"""tcp"""'}, {}), "(self.remote_port, 'tcp')", False, 'import socket\n'), ((69, 19, 69, 57), 'socket.gethostbyaddr', 'socket.gethostbyaddr', ({(69, 40, 69, 56): 'self.remote_addr'}, {}), '(self.remote_addr)', False, 'import socket\n'), ((81, 36, 81, 57), 'ctypes.byref', 'ctypes.byref', ({(81, 49, 81, 56): 'closing'}, {}), '(closing)', False, 'import ctypes\n'), ((421, 45, 421, 63), 'ctypes.byref', 'ctypes.byref', ({(421, 58, 421, 62): 'size'}, {}), '(size)', False, 'import ctypes\n'), ((433, 45, 433, 63), 'ctypes.byref', 'ctypes.byref', ({(433, 58, 433, 62): 'size'}, {}), '(size)', False, 'import ctypes\n'), ((417, 47, 417, 65), 'ctypes.byref', 'ctypes.byref', ({(417, 60, 417, 64): 'size'}, {}), '(size)', False, 'import ctypes\n'), ((429, 47, 429, 65), 'ctypes.byref', 'ctypes.byref', ({(429, 60, 429, 64): 'size'}, {}), '(size)', False, 'import ctypes\n')]