Skip to content

Commit

Permalink
Merge pull request #28 from bung87/minor-improve-windows
Browse files Browse the repository at this point in the history
minor improve windows version
  • Loading branch information
juancarlospaco authored Jan 12, 2024
2 parents 79c3580 + 5ade984 commit 65098a9
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 70 deletions.
81 changes: 13 additions & 68 deletions src/psutil/psutil_windows.nim
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
{.deadCodeElim: on.}

import math
import sequtils
import strformat
Expand Down Expand Up @@ -57,7 +55,6 @@ proc psutil_get_drive_type*(drive_type: UINT): string =
else: "?"

proc psutil_get_drive_type*(drive: string): string =

var drive_type = GetDriveType(drive)
case drive_type
of DRIVE_FIXED: "fixed"
Expand All @@ -72,40 +69,25 @@ proc psutil_get_drive_type*(drive: string): string =
proc getnativearch*(): int =
## Get the native architecture of the system we are running on
var pGetNativeSystemInfo: SYSTEM_INFO
var nativeArch = PROCESS_ARCH_UNKNOWN
result = PROCESS_ARCH_X86

GetNativeSystemInfo(pGetNativeSystemInfo.addr)

if pGetNativeSystemInfo.isNil:
raiseError()


case pGetNativeSystemInfo.union1.struct1.wProcessorArchitecture
result = case pGetNativeSystemInfo.union1.struct1.wProcessorArchitecture
of PROCESSOR_ARCHITECTURE_AMD64:
## 64 bit (x64)
# dwNativeArch = PROCESSOR_ARCHITECTURE_AMD64
nativeArch = PROCESS_ARCH_X64

PROCESS_ARCH_X64
of PROCESSOR_ARCHITECTURE_IA64:
# dwNativeArch = PROCESSOR_ARCHITECTURE_IA64
nativeArch = PROCESS_ARCH_X64

PROCESS_ARCH_X64
of PROCESSOR_ARCHITECTURE_INTEL:
# 32 bit (x86)
# dwNativeArch = PROCESSOR_ARCHITECTURE_INTEL
nativeArch = PROCESS_ARCH_X64

PROCESS_ARCH_X64
else:
# dwNativeArch = PROCESSOR_ARCHITECTURE_UNKNOWN
nativeArch = PROCESS_ARCH_UNKNOWN

return nativeArch

PROCESS_ARCH_UNKNOWN

proc pids*(): seq[int] =
## Returns a list of PIDs currently running on the system.
result = newSeq[int]()

var procArray: seq[DWORD]
var procArrayLen = 0
# Stores the byte size of the returned array from enumprocesses
Expand All @@ -119,7 +101,6 @@ proc pids*(): seq[int] =
DWORD(procArrayLen * sizeof(DWORD)),
addr enumReturnSz) == 0:
raiseError()
return result

# The number of elements is the returned size / size of each element
let numberOfReturnedPIDs = int(int(enumReturnSz) / sizeof(DWORD))
Expand Down Expand Up @@ -169,12 +150,10 @@ proc pid_path*(pid: int): string =
result = $filename

proc pid_paths*(pids: seq[int]): seq[string] =

for pid in pids:
result.add(pid_path(pid))

proc try_pid_path*(pid: int): string =

var processHandle: HANDLE
var filename: wstring = newWString(MAX_PATH)
var dwSize = MAX_PATH
Expand All @@ -193,13 +172,11 @@ proc try_pid_path*(pid: int): string =


proc try_pid_paths*(pids: seq[int]): seq[string] =

## Function to return the paths of the exes (sequence of strings) of the running pids.
for pid in pids:
result.add(try_pid_path(pid))

proc pid_parent*(pid: int): int =

var h: HANDLE
var pe: PROCESSENTRY32
var ppid = DWORD(0)
Expand All @@ -215,45 +192,28 @@ proc pid_parent*(pid: int): int =
return cast[int](ppid)

proc pid_parents*(pids: seq[int]): seq[int] =

var ret: seq[int]
for pid in pids:
ret.add(pid_parent(pid))

return ret
result.add(pid_parent(pid))

proc pids_with_names*(): (seq[int], seq[string]) =

## Function for returning tuple of pids and names

var pids_seq = pids()
var names_seq = pid_names(pids_seq)

return (pids_seq, names_seq)


proc pid_arch*(pid: int): int =

## function for getting the architecture of the pid running
var bIsWow64: BOOL
var nativeArch = static PROCESS_ARCH_UNKNOWN
var hProcess: HANDLE
# var pIsWow64Process: ISWOW64PROCESS
var dwPid = DWORD(pid)
# now we must default to an unknown architecture as the process may be either x86/x64 and we may not have the rights to open it
result = PROCESS_ARCH_UNKNOWN

## grab the native systems architecture the first time we use this function we use this funciton.
if nativeArch == PROCESS_ARCH_UNKNOWN:
nativeArch = getnativearch()


var nativeArch = getnativearch()
result = PROCESS_ARCH_X86
hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPid)
defer: CloseHandle(hProcess)
if hProcess == cast[HANDLE](-1):
if hProcess == HANDLE(0):
hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, dwPid)
if hProcess == cast[HANDLE](-1):
raiseError()
if hProcess == HANDLE(0):
return

if IsWow64Process(hProcess, bIsWow64.addr) == FALSE:
return
Expand All @@ -264,7 +224,6 @@ proc pid_arch*(pid: int): int =
result = nativeArch

proc pid_user*(pid: int): string =

## Attempt to get the username associated with the given pid.
var hProcess: HANDLE
var hToken: HANDLE
Expand Down Expand Up @@ -305,14 +264,11 @@ proc pid_user*(pid: int): string =
result = $wcUser

proc pid_users*(pids: seq[int]): seq[string] =

## Function for getting a sequence of users
for pid in pids:
result.add(pid_user(pid))


proc try_pid_user*(pid: int): string =

## Attempt to get the username associated with the given pid.
var hProcess: HANDLE
var hToken: HANDLE
Expand All @@ -325,7 +281,6 @@ proc try_pid_user*(pid: int): string =
var wcUser: wstring = newWstring(512)
var wcDomain: wstring = newWstring(512)


hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPid)
defer: CloseHandle(hProcess)
if hProcess == DWORD(0):
Expand Down Expand Up @@ -355,13 +310,11 @@ proc try_pid_user*(pid: int): string =


proc try_pid_users*(pids: seq[int]): seq[string] =

## Function for getting users of specified pids
for pid in pids:
result.add(try_pid_user(pid))

proc pid_domain*(pid: int): string =

## Attempt to get the domain associated with the given pid.
var hProcess: HANDLE
var hToken: HANDLE
Expand All @@ -374,7 +327,6 @@ proc pid_domain*(pid: int): string =
var wcUser: wstring = newWstring(512)
var wcDomain: wstring = newWstring(512)


hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPid)
defer: CloseHandle(hProcess)
if hProcess == DWORD(0):
Expand Down Expand Up @@ -403,7 +355,6 @@ proc pid_domain*(pid: int): string =


proc pid_domain_user*(pid: int): (string, string) =

## Attempt to get the domain and username associated with the given pid.
var hProcess: HANDLE
var hToken: HANDLE
Expand All @@ -416,7 +367,6 @@ proc pid_domain_user*(pid: int): (string, string) =
var wcUser: wstring = newWstring(512)
var wcDomain: wstring = newWstring(512)


hProcess = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPid)
defer: CloseHandle(hProcess)
if hProcess == DWORD(0):
Expand All @@ -427,7 +377,7 @@ proc pid_domain_user*(pid: int): (string, string) =

defer: CloseHandle(hToken)

if hToken == cast[HANDLE](-1) or hToken == cast[HANDLE](NULL):
if hToken == HANDLE(0):
raiseError()

## Get required buffer size and allocate the TOKEN_USER buffer
Expand Down Expand Up @@ -905,15 +855,10 @@ proc process_exists*(processName: string): bool =
return exists

proc pid_exists*(pid: int): bool =

var p = OpenProcess(SYNCHRONIZE, FALSE, DWORD(pid));
var r = WaitForSingleObject(p, 0);
CloseHandle(p);
return r == WAIT_TIMEOUT



proc pid_cmdline*(pid: int): string =
raise newException(Exception, "Function is unimplemented!")


3 changes: 1 addition & 2 deletions tests/test_windows.nim
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@ import ../src/psutil
when defined(windows):
import ../src/psutil/psutil_windows
import winim
echo getnativearch()
# echo_proc pid_exists(77724)
echo "arch",getnativearch()
let ids = pids()
echo ids
# echo pids_with_names()
Expand Down

0 comments on commit 65098a9

Please sign in to comment.