codekingpro/portable-devtools

1# -*- coding: utf-8 -*-2# Copyright (C) 2016  Fabio Falcinelli, Maximilian Hils3#4# This program is free software: you can redistribute it and/or modify5# it under the terms of the GNU Lesser General Public License as published by6# the Free Software Foundation, either version 3 of the License, or7# (at your option) any later version.8#9# This program is distributed in the hope that it will be useful,10# but WITHOUT ANY WARRANTY; without even the implied warranty of11# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the12# GNU Lesser General Public License for more details.13#14# You should have received a copy of the GNU Lesser General Public License15# along with this program.  If not, see <http://www.gnu.org/licenses/>.16import subprocess17import sys18from ctypes import byref, c_uint64, c_uint, c_char, c_char_p19 20from pydivert import windivert_dll21from pydivert.consts import Layer, Direction, Flag22from pydivert.packet import Packet23from pydivert.util import PY224 25DEFAULT_PACKET_BUFFER_SIZE = 150026 27 28class WinDivert(object):29    """30    A WinDivert handle that can be used to capture packets.31    The main methods are `.open()`, `.recv()`, `.send()` and `.close()`.32 33    Use it like so::34 35        with pydivert.WinDivert() as w:36            for packet in w:37                print(packet)38                w.send(packet)39 40    """41 42    def __init__(self, filter="true", layer=Layer.NETWORK, priority=0, flags=Flag.DEFAULT):43        self._handle = None44        self._filter = filter.encode()45        self._layer = layer46        self._priority = priority47        self._flags = flags48 49    def __repr__(self):50        return '<WinDivert state="{}" filter="{}" layer="{}" priority="{}" flags="{}" />'.format(51            "open" if self._handle is not None else "closed",52            self._filter.decode(),53            self._layer,54            self._priority,55            self._flags56        )57 58    def __enter__(self):59        self.open()60        return self61 62    def __exit__(self, *args):63        self.close()64 65    def __iter__(self):66        return self67 68    def __next__(self):69        return self.recv()70 71    if sys.version_info < (3, 0):72        next = __next__73 74    @staticmethod75    def register():76        """77        An utility method to register the service the first time.78        It is usually not required to call this function, as WinDivert will register itself when opening a handle.79        """80        with WinDivert("false"):81            pass82 83    @staticmethod84    def is_registered():85        """86        Check if the WinDivert service is currently installed on the system.87        """88        return subprocess.call("sc query WinDivert1.3", stdout=subprocess.PIPE,89                               stderr=subprocess.PIPE) == 090 91    @staticmethod92    def unregister():93        """94        Unregisters the WinDivert service.95        This function only requests a service stop, which may not be processed immediately if there are still open96        handles.97        """98        subprocess.check_call("sc stop WinDivert1.3", stdout=subprocess.PIPE,99                              stderr=subprocess.PIPE)100 101    @staticmethod102    def check_filter(filter, layer=Layer.NETWORK):103        """104        Checks if the given packet filter string is valid with respect to the filter language.105 106        The remapped function is WinDivertHelperCheckFilter::107 108            BOOL WinDivertHelperCheckFilter(109                __in const char *filter,110                __in WINDIVERT_LAYER layer,111                __out_opt const char **errorStr,112                __out_opt UINT *errorPos113            );114 115        See: https://reqrypt.org/windivert-doc.html#divert_helper_check_filter116 117        :return: A tuple (res, pos, msg) with check result in 'res' human readable description of the error in 'msg' and the error's position in 'pos'.118        """119        res, pos, msg = False, c_uint(), c_char_p()120        try:121            res = windivert_dll.WinDivertHelperCheckFilter(filter.encode(), layer, byref(msg), byref(pos))122        except OSError:123            pass124        return res, pos.value, msg.value.decode()125 126    def open(self):127        """128        Opens a WinDivert handle for the given filter.129        Unless otherwise specified by flags, any packet that matches the filter will be diverted to the handle.130        Diverted packets can be read by the application with receive().131 132        The remapped function is WinDivertOpen::133 134            HANDLE WinDivertOpen(135                __in const char *filter,136                __in WINDIVERT_LAYER layer,137                __in INT16 priority,138                __in UINT64 flags139            );140 141        For more info on the C call visit: http://reqrypt.org/windivert-doc.html#divert_open142        """143        if self.is_open:144            raise RuntimeError("WinDivert handle is already open.")145        self._handle = windivert_dll.WinDivertOpen(self._filter, self._layer, self._priority,146                                                   self._flags)147 148    @property149    def is_open(self):150        """151        Indicates if there is currently an open handle.152        """153        return bool(self._handle)154 155    def close(self):156        """157        Closes the handle opened by open().158 159        The remapped function is WinDivertClose::160 161            BOOL WinDivertClose(162                __in HANDLE handle163            );164 165        For more info on the C call visit: http://reqrypt.org/windivert-doc.html#divert_close166        """167        if not self.is_open:168            raise RuntimeError("WinDivert handle is not open.")169        windivert_dll.WinDivertClose(self._handle)170        self._handle = None171 172    def recv(self, bufsize=DEFAULT_PACKET_BUFFER_SIZE):173        """174        Receives a diverted packet that matched the filter.175 176        The remapped function is WinDivertRecv::177 178            BOOL WinDivertRecv(179                __in HANDLE handle,180                __out PVOID pPacket,181                __in UINT packetLen,182                __out_opt PWINDIVERT_ADDRESS pAddr,183                __out_opt UINT *recvLen184            );185 186        For more info on the C call visit: http://reqrypt.org/windivert-doc.html#divert_recv187 188        :return: The return value is a `pydivert.Packet`.189        """190        if self._handle is None:191            raise RuntimeError("WinDivert handle is not open")192 193        packet = bytearray(bufsize)194        packet_ = (c_char * bufsize).from_buffer(packet)195        address = windivert_dll.WinDivertAddress()196        recv_len = c_uint(0)197        windivert_dll.WinDivertRecv(self._handle, packet_, bufsize, byref(address), byref(recv_len))198        return Packet(199            memoryview(packet)[:recv_len.value],200            (address.IfIdx, address.SubIfIdx),201            Direction(address.Direction)202        )203 204    def send(self, packet, recalculate_checksum=True):205        """206        Injects a packet into the network stack.207        Recalculates the checksum before sending unless recalculate_checksum=False is passed.208 209        The injected packet may be one received from recv(), or a modified version, or a completely new packet.210        Injected packets can be captured and diverted again by other WinDivert handles with lower priorities.211 212        The remapped function is WinDivertSend::213 214            BOOL WinDivertSend(215                __in HANDLE handle,216                __in PVOID pPacket,217                __in UINT packetLen,218                __in PWINDIVERT_ADDRESS pAddr,219                __out_opt UINT *sendLen220            );221 222        For more info on the C call visit: http://reqrypt.org/windivert-doc.html#divert_send223 224        :return: The return value is the number of bytes actually sent.225        """226        if recalculate_checksum:227            packet.recalculate_checksums()228 229        send_len = c_uint(0)230        if PY2:231            # .from_buffer(memoryview) does not work on PY2232            buff = bytearray(packet.raw)233        else:234            buff = packet.raw235        buff = (c_char * len(packet.raw)).from_buffer(buff)236        windivert_dll.WinDivertSend(self._handle, buff, len(packet.raw), byref(packet.wd_addr),237                                    byref(send_len))238        return send_len239 240    def get_param(self, name):241        """242        Get a WinDivert parameter. See pydivert.Param for the list of parameters.243 244        The remapped function is WinDivertGetParam::245 246            BOOL WinDivertGetParam(247                __in HANDLE handle,248                __in WINDIVERT_PARAM param,249                __out UINT64 *pValue250            );251 252        For more info on the C call visit: http://reqrypt.org/windivert-doc.html#divert_get_param253 254        :return: The parameter value.255        """256        value = c_uint64(0)257        windivert_dll.WinDivertGetParam(self._handle, name, byref(value))258        return value.value259 260    def set_param(self, name, value):261        """262        Set a WinDivert parameter. See pydivert.Param for the list of parameters.263 264        The remapped function is DivertSetParam::265 266            BOOL WinDivertSetParam(267                __in HANDLE handle,268                __in WINDIVERT_PARAM param,269                __in UINT64 value270            );271 272        For more info on the C call visit: http://reqrypt.org/windivert-doc.html#divert_set_param273        """274        return windivert_dll.WinDivertSetParam(self._handle, name, value)275