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/>.16"""17pydivert bundles the WinDivert binaries from18https://reqrypt.org/download/WinDivert-1.3.0-WDDK.zip19"""20import functools21import os22import platform23import sys24from ctypes import (25    POINTER, GetLastError, WinError, c_uint, c_void_p, c_uint32, c_char_p, ARRAY, c_uint64, c_int16, c_int, WinDLL,26    c_uint8, windll)27from ctypes.wintypes import HANDLE28 29from .structs import WinDivertAddress30 31ERROR_IO_PENDING = 99732 33here = os.path.abspath(os.path.dirname(__file__))34 35if platform.architecture()[0] == "64bit":36    DLL_PATH = os.path.join(here, "WinDivert64.dll")37else:38    DLL_PATH = os.path.join(here, "WinDivert32.dll")39 40 41def raise_on_error(f):42    """43    This decorator throws a WinError whenever GetLastError() returns an error.44    As as special case, ERROR_IO_PENDING is ignored.45    """46 47    @functools.wraps(f)48    def wrapper(*args, **kwargs):49        result = f(*args, **kwargs)50        retcode = GetLastError()51        if retcode and retcode != ERROR_IO_PENDING:52            err = WinError(code=retcode)53            windll.kernel32.SetLastError(0)  # clear error code so that we don't raise twice.54            raise err55        return result56 57    return wrapper58 59 60WINDIVERT_FUNCTIONS = {61    "WinDivertHelperParsePacket": [HANDLE, c_uint, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p, c_void_p,62                                   c_void_p, POINTER(c_uint)],63    "WinDivertHelperParseIPv4Address": [c_char_p, POINTER(c_uint32)],64    "WinDivertHelperParseIPv6Address": [c_char_p, POINTER(ARRAY(c_uint8, 16))],65    "WinDivertHelperCalcChecksums": [c_void_p, c_uint, c_uint64],66    "WinDivertHelperCheckFilter": [c_char_p, c_int, POINTER(c_char_p), POINTER(c_uint)],67    "WinDivertHelperEvalFilter": [c_char_p, c_int, c_void_p, c_uint, c_void_p],68    "WinDivertOpen": [c_char_p, c_int, c_int16, c_uint64],69    "WinDivertRecv": [HANDLE, c_void_p, c_uint, c_void_p, c_void_p],70    "WinDivertSend": [HANDLE, c_void_p, c_uint, c_void_p, c_void_p],71    "WinDivertRecvEx": [HANDLE, c_void_p, c_uint, c_uint64, c_void_p, c_void_p, c_void_p],72    "WinDivertSendEx": [HANDLE, c_void_p, c_uint, c_uint64, c_void_p, c_void_p, c_void_p],73    "WinDivertClose": [HANDLE],74    "WinDivertGetParam": [HANDLE, c_int, POINTER(c_uint64)],75    "WinDivertSetParam": [HANDLE, c_int, c_uint64],76}77 78_instance = None79 80 81def instance():82    global _instance83    if _instance is None:84        _instance = WinDLL(DLL_PATH)85        for funcname, argtypes in WINDIVERT_FUNCTIONS.items():86            func = getattr(_instance, funcname)87            func.argtypes = argtypes88    return _instance89 90 91# Dark magic happens below.92# On init, windivert_dll.WinDivertOpen is a proxy function that loads the DLL on the first invocation93# and then replaces all existing proxy function with direct handles to the DLL's functions.94 95 96_module = sys.modules[__name__]97 98 99def _init():100    """101    Lazy-load DLL, replace proxy functions with actual ones.102    """103    i = instance()104    for funcname in WINDIVERT_FUNCTIONS:105        func = getattr(i, funcname)106        func = raise_on_error(func)107        setattr(_module, funcname, func)108 109 110def _mkprox(funcname):111    """112    Make lazy-init proxy function.113    """114 115    def prox(*args, **kwargs):116        _init()117        return getattr(_module, funcname)(*args, **kwargs)118 119    return prox120 121 122for funcname in WINDIVERT_FUNCTIONS:123    setattr(_module, funcname, _mkprox(funcname))124