# -*- coding: utf-8 -*-
"""
billiard._win
~~~~~~~~~~~~~
Windows utilities to terminate process groups.
"""
from __future__ import absolute_import
import os
# psutil is painfully slow in win32. So to avoid adding big
# dependencies like pywin32 a ctypes based solution is preferred
# Code based on the winappdbg project http://winappdbg.sourceforge.net/
# (BSD License)
from ctypes import (
byref, sizeof, windll,
Structure, WinError, POINTER,
c_size_t, c_char, c_void_p,
)
from ctypes.wintypes import DWORD, LONG
ERROR_NO_MORE_FILES = 18
INVALID_HANDLE_VALUE = c_void_p(-1).value
class PROCESSENTRY32(Structure):
_fields_ = [
('dwSize', DWORD),
('cntUsage', DWORD),
('th32ProcessID', DWORD),
('th32DefaultHeapID', c_size_t),
('th32ModuleID', DWORD),
('cntThreads', DWORD),
('th32ParentProcessID', DWORD),
('pcPriClassBase', LONG),
('dwFlags', DWORD),
('szExeFile', c_char * 260),
]
LPPROCESSENTRY32 = POINTER(PROCESSENTRY32)
def CreateToolhelp32Snapshot(dwFlags=2, th32ProcessID=0):
hSnapshot = windll.kernel32.CreateToolhelp32Snapshot(dwFlags,
th32ProcessID)
if hSnapshot == INVALID_HANDLE_VALUE:
raise WinError()
return hSnapshot
def Process32First(hSnapshot, pe=None):
return _Process32n(windll.kernel32.Process32First, hSnapshot, pe)
def Process32Next(hSnapshot, pe=None):
return _Process32n(windll.kernel32.Process32Next, hSnapshot, pe)
def _Process32n(fun, hSnapshot, pe=None):
if pe is None:
pe = PROCESSENTRY32()
pe.dwSize = sizeof(PROCESSENTRY32)
success = fun(hSnapshot, byref(pe))
if not success:
if windll.kernel32.GetLastError() == ERROR_NO_MORE_FILES:
return
raise WinError()
return pe
def get_all_processes_pids():
"""Return a dictionary with all processes pids as keys and their
parents as value. Ignore processes with no parents.
"""
h = CreateToolhelp32Snapshot()
parents = {}
pe = Process32First(h)
while pe:
if pe.th32ParentProcessID:
parents[pe.th32ProcessID] = pe.th32ParentProcessID
pe = Process32Next(h, pe)
return parents
def get_processtree_pids(pid, include_parent=True):
"""Return a list with all the pids of a process tree"""
parents = get_all_processes_pids()
all_pids = list(parents.keys())
pids = set([pid])
while 1:
pids_new = pids.copy()
for _pid in all_pids:
if parents[_pid] in pids:
pids_new.add(_pid)
if pids_new == pids:
break
pids = pids_new.copy()
if not include_parent:
pids.remove(pid)
return list(pids)
def kill_processtree(pid, signum):
"""Kill a process and all its descendants"""
family_pids = get_processtree_pids(pid)
for _pid in family_pids:
os.kill(_pid, signum)