1# Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Common objects shared by __init__.py and _ps*.py modules.
6
7Note: this module is imported by setup.py, so it should not import
8psutil or third-party modules.
9"""
10
11import collections
12import functools
13import os
14import socket
15import stat
16import sys
17import threading
18import warnings
19from socket import AF_INET
20from socket import SOCK_DGRAM
21from socket import SOCK_STREAM
22
23try:
24 from socket import AF_INET6
25except ImportError:
26 AF_INET6 = None
27try:
28 from socket import AF_UNIX
29except ImportError:
30 AF_UNIX = None
31
32
33PSUTIL_DEBUG = bool(os.getenv('PSUTIL_DEBUG'))
34PSUTIL_TESTING = bool(os.getenv('PSUTIL_TESTING'))
35_DEFAULT = object()
36
37# fmt: off
38__all__ = [
39 # OS constants
40 'FREEBSD', 'BSD', 'LINUX', 'NETBSD', 'OPENBSD', 'MACOS', 'POSIX',
41 'SUNOS', 'WINDOWS',
42 # other constants
43 'ENCODING', 'ENCODING_ERRS', 'AF_INET6',
44 # utility functions
45 'conn_tmap', 'isfile_strict',
46 'parse_environ_block', 'path_exists_strict', 'usage_percent',
47 'supports_ipv6', 'sockfam_to_enum', 'socktype_to_enum', "wrap_numbers",
48 'open_text', 'open_binary', 'cat', 'bcat',
49 'bytes2human', 'conn_to_ntuple', 'debug', 'warn',
50 # shell utils
51 'hilite', 'term_supports_colors', 'print_color',
52]
53# fmt: on
54
55
56# ===================================================================
57# --- OS constants
58# ===================================================================
59
60
61POSIX = os.name == "posix"
62WINDOWS = os.name == "nt"
63LINUX = sys.platform.startswith(("linux", "android"))
64MACOS = sys.platform.startswith("darwin")
65FREEBSD = sys.platform.startswith(("freebsd", "midnightbsd"))
66OPENBSD = sys.platform.startswith("openbsd")
67NETBSD = sys.platform.startswith("netbsd")
68BSD = FREEBSD or OPENBSD or NETBSD
69SUNOS = sys.platform.startswith(("sunos", "solaris"))
70AIX = sys.platform.startswith("aix")
71
72ENCODING = sys.getfilesystemencoding()
73ENCODING_ERRS = sys.getfilesystemencodeerrors()
74
75
76# ===================================================================
77# --- Process.net_connections() 'kind' parameter mapping
78# ===================================================================
79
80
81conn_tmap = {
82 "all": ([AF_INET, AF_INET6, AF_UNIX], [SOCK_STREAM, SOCK_DGRAM]),
83 "tcp": ([AF_INET, AF_INET6], [SOCK_STREAM]),
84 "tcp4": ([AF_INET], [SOCK_STREAM]),
85 "udp": ([AF_INET, AF_INET6], [SOCK_DGRAM]),
86 "udp4": ([AF_INET], [SOCK_DGRAM]),
87 "inet": ([AF_INET, AF_INET6], [SOCK_STREAM, SOCK_DGRAM]),
88 "inet4": ([AF_INET], [SOCK_STREAM, SOCK_DGRAM]),
89 "inet6": ([AF_INET6], [SOCK_STREAM, SOCK_DGRAM]),
90}
91
92if AF_INET6 is not None:
93 conn_tmap.update({
94 "tcp6": ([AF_INET6], [SOCK_STREAM]),
95 "udp6": ([AF_INET6], [SOCK_DGRAM]),
96 })
97
98if AF_UNIX is not None and not SUNOS:
99 conn_tmap.update({"unix": ([AF_UNIX], [SOCK_STREAM, SOCK_DGRAM])})
100
101
102# =====================================================================
103# --- Exceptions
104# =====================================================================
105
106
107class Error(Exception):
108 """Base exception class. All other psutil exceptions inherit
109 from this one.
110 """
111
112 __module__ = 'psutil'
113
114 def _infodict(self, attrs):
115 info = {}
116 for name in attrs:
117 value = getattr(self, name, None)
118 if value or (name == "pid" and value == 0):
119 info[name] = value
120 return info
121
122 def __str__(self):
123 # invoked on `raise Error`
124 if info := self._infodict(("pid", "ppid", "name")):
125 details = "({})".format(
126 ", ".join([f"{k}={v!r}" for k, v in info.items()])
127 )
128 else:
129 details = None
130 return " ".join([x for x in (getattr(self, "msg", ""), details) if x])
131
132 def __repr__(self):
133 # invoked on `repr(Error)`
134 info = self._infodict(("pid", "ppid", "name", "seconds", "msg"))
135 details = ", ".join([f"{k}={v!r}" for k, v in info.items()])
136 return f"psutil.{self.__class__.__name__}({details})"
137
138
139class NoSuchProcess(Error):
140 """Exception raised when a process with a certain PID doesn't
141 or no longer exists.
142 """
143
144 __module__ = 'psutil'
145
146 def __init__(self, pid, name=None, msg=None):
147 Error.__init__(self)
148 self.pid = pid
149 self.name = name
150 self.msg = msg or "process no longer exists"
151
152 def __reduce__(self):
153 return (self.__class__, (self.pid, self.name, self.msg))
154
155
156class ZombieProcess(NoSuchProcess):
157 """Exception raised when querying a zombie process. This is
158 raised on macOS, BSD and Solaris only, and not always: depending
159 on the query the OS may be able to succeed anyway.
160 On Linux all zombie processes are querable (hence this is never
161 raised). Windows doesn't have zombie processes.
162 """
163
164 __module__ = 'psutil'
165
166 def __init__(self, pid, name=None, ppid=None, msg=None):
167 NoSuchProcess.__init__(self, pid, name, msg)
168 self.ppid = ppid
169 self.msg = msg or "PID still exists but it's a zombie"
170
171 def __reduce__(self):
172 return (self.__class__, (self.pid, self.name, self.ppid, self.msg))
173
174
175class AccessDenied(Error):
176 """Exception raised when permission to perform an action is denied."""
177
178 __module__ = 'psutil'
179
180 def __init__(self, pid=None, name=None, msg=None):
181 Error.__init__(self)
182 self.pid = pid
183 self.name = name
184 self.msg = msg or ""
185
186 def __reduce__(self):
187 return (self.__class__, (self.pid, self.name, self.msg))
188
189
190class TimeoutExpired(Error):
191 """Raised on Process.wait(timeout) if timeout expires and process
192 is still alive.
193 """
194
195 __module__ = 'psutil'
196
197 def __init__(self, seconds, pid=None, name=None):
198 Error.__init__(self)
199 self.seconds = seconds
200 self.pid = pid
201 self.name = name
202 self.msg = f"timeout after {seconds} seconds"
203
204 def __reduce__(self):
205 return (self.__class__, (self.seconds, self.pid, self.name))
206
207
208# ===================================================================
209# --- utils
210# ===================================================================
211
212
213def usage_percent(used, total, round_=None):
214 """Calculate percentage usage of 'used' against 'total'."""
215 try:
216 ret = (float(used) / total) * 100
217 except ZeroDivisionError:
218 return 0.0
219 else:
220 if round_ is not None:
221 ret = round(ret, round_)
222 return ret
223
224
225def cache_activate(proc):
226 """Make `memoize_when_activated` methods of *proc* start caching."""
227 proc._cache = {}
228
229
230def cache_deactivate(proc):
231 """Stop caching and discard what was collected."""
232 try:
233 del proc._cache
234 except AttributeError:
235 pass
236
237
238def memoize_when_activated(fun):
239 """A memoize decorator which is disabled by default. It can be
240 activated and deactivated on request.
241 For efficiency reasons it can be used only against class methods
242 accepting no arguments.
243
244 >>> class Foo:
245 ... @memoize_when_activated
246 ... def foo()
247 ... print(1)
248 ...
249 >>> f = Foo()
250 >>> # deactivated (default)
251 >>> foo()
252 1
253 >>> foo()
254 1
255 >>>
256 >>> # activated
257 >>> cache_activate(f)
258 >>> foo()
259 1
260 >>> foo()
261 >>> foo()
262 >>>
263 """
264
265 @functools.wraps(fun)
266 def wrapper(self):
267 try:
268 # case 1: we previously entered oneshot() ctx
269 ret = self._cache[fun]
270 except AttributeError:
271 # case 2: we never entered oneshot() ctx
272 try:
273 return fun(self)
274 except Exception as err:
275 raise err from None
276 except KeyError:
277 # case 3: we entered oneshot() ctx but there's no cache
278 # for this entry yet
279 try:
280 ret = fun(self)
281 except Exception as err:
282 raise err from None
283 try:
284 self._cache[fun] = ret
285 except AttributeError:
286 # multi-threading race condition, see:
287 # https://github.com/giampaolo/psutil/issues/1948
288 pass
289 return ret
290
291 return wrapper
292
293
294def isfile_strict(path):
295 """Same as os.path.isfile() but does not swallow EACCES / EPERM
296 exceptions, see:
297 http://mail.python.org/pipermail/python-dev/2012-June/120787.html.
298 """
299 try:
300 st = os.stat(path)
301 except PermissionError:
302 raise
303 except OSError:
304 return False
305 else:
306 return stat.S_ISREG(st.st_mode)
307
308
309def path_exists_strict(path):
310 """Same as os.path.exists() but does not swallow EACCES / EPERM
311 exceptions. See:
312 http://mail.python.org/pipermail/python-dev/2012-June/120787.html.
313 """
314 try:
315 os.stat(path)
316 except PermissionError:
317 raise
318 except OSError:
319 return False
320 else:
321 return True
322
323
324def supports_ipv6():
325 """Return True if IPv6 is supported on this platform."""
326 if not socket.has_ipv6 or AF_INET6 is None:
327 return False
328 try:
329 with socket.socket(AF_INET6, socket.SOCK_STREAM) as sock:
330 sock.bind(("::1", 0))
331 return True
332 except OSError:
333 return False
334
335
336def parse_environ_block(data):
337 """Parse a C environ block of environment variables into a dictionary."""
338 # The block is usually raw data from the target process. It might contain
339 # trailing garbage and lines that do not look like assignments.
340 ret = {}
341 pos = 0
342
343 # localize global variable to speed up access.
344 WINDOWS_ = WINDOWS
345 while True:
346 next_pos = data.find("\0", pos)
347 # nul byte at the beginning or double nul byte means finish
348 if next_pos <= pos:
349 break
350 # there might not be an equals sign
351 equal_pos = data.find("=", pos, next_pos)
352 if equal_pos > pos:
353 key = data[pos:equal_pos]
354 value = data[equal_pos + 1 : next_pos]
355 # Windows expects environment variables to be uppercase only
356 if WINDOWS_:
357 key = key.upper()
358 ret[key] = value
359 pos = next_pos + 1
360
361 return ret
362
363
364def sockfam_to_enum(num):
365 """Convert a numeric socket family value to an IntEnum member.
366 If it's not a known member, return the numeric value itself.
367 """
368 try:
369 return socket.AddressFamily(num)
370 except ValueError:
371 return num
372
373
374def socktype_to_enum(num):
375 """Convert a numeric socket type value to an IntEnum member.
376 If it's not a known member, return the numeric value itself.
377 """
378 try:
379 return socket.SocketKind(num)
380 except ValueError:
381 return num
382
383
384def conn_to_ntuple(fd, fam, type_, laddr, raddr, status, status_map, pid=None):
385 """Convert a raw connection tuple to a proper ntuple."""
386 from . import _ntuples as ntp
387 from ._enums import ConnectionStatus
388
389 if fam in {socket.AF_INET, AF_INET6}:
390 if laddr:
391 laddr = ntp.addr(*laddr)
392 if raddr:
393 raddr = ntp.addr(*raddr)
394 if type_ == socket.SOCK_STREAM and fam in {AF_INET, AF_INET6}:
395 status = status_map.get(status, ConnectionStatus.CONN_NONE)
396 else:
397 status = ConnectionStatus.CONN_NONE # ignore whatever C returned to us
398 fam = sockfam_to_enum(fam)
399 type_ = socktype_to_enum(type_)
400 if pid is None:
401 return ntp.pconn(fd, fam, type_, laddr, raddr, status)
402 else:
403 return ntp.sconn(fd, fam, type_, laddr, raddr, status, pid)
404
405
406def broadcast_addr(addr):
407 """Given the address ntuple returned by ``net_if_addrs()``
408 calculates the broadcast address. Returns None for a single-host
409 network (/32 or /128), which has no broadcast address.
410 """
411 import ipaddress
412
413 if not addr.address or not addr.netmask:
414 return None
415 if addr.family == socket.AF_INET:
416 net = ipaddress.IPv4Network(
417 f"{addr.address}/{addr.netmask}", strict=False
418 )
419 elif addr.family == socket.AF_INET6:
420 net = ipaddress.IPv6Network(
421 f"{addr.address}/{addr.netmask}", strict=False
422 )
423 else:
424 return None
425 if net.prefixlen == net.max_prefixlen:
426 return None
427 return str(net.broadcast_address)
428
429
430class _WrapNumbers:
431 """Watches numbers so that they don't overflow and wrap
432 (reset to zero).
433 """
434
435 def __init__(self):
436 self.lock = threading.Lock()
437 self.cache = {}
438 self.reminders = {}
439 self.reminder_keys = {}
440
441 def _add_dict(self, input_dict, name):
442 assert name not in self.cache
443 assert name not in self.reminders
444 assert name not in self.reminder_keys
445 self.cache[name] = input_dict
446 self.reminders[name] = collections.defaultdict(int)
447 self.reminder_keys[name] = collections.defaultdict(set)
448
449 def _remove_dead_reminders(self, input_dict, name):
450 """In case the number of keys changed between calls (e.g. a
451 disk disappears) this removes the entry from self.reminders.
452 """
453 old_dict = self.cache[name]
454 gone_keys = set(old_dict) - set(input_dict)
455 for gone_key in gone_keys:
456 for remkey in self.reminder_keys[name][gone_key]:
457 del self.reminders[name][remkey]
458 del self.reminder_keys[name][gone_key]
459
460 def run(self, input_dict, name):
461 """Cache dict and sum numbers which overflow and wrap.
462 Return an updated copy of `input_dict`.
463 """
464 if name not in self.cache:
465 # This was the first call.
466 self._add_dict(input_dict, name)
467 return input_dict
468
469 self._remove_dead_reminders(input_dict, name)
470
471 old_dict = self.cache[name]
472 new_dict = {}
473 for key in input_dict:
474 input_tuple = input_dict[key]
475 try:
476 old_tuple = old_dict[key]
477 except KeyError:
478 # The input dict has a new key (e.g. a new disk or NIC)
479 # which didn't exist in the previous call.
480 new_dict[key] = input_tuple
481 continue
482
483 bits = []
484 for i in range(len(input_tuple)):
485 input_value = input_tuple[i]
486 old_value = old_tuple[i]
487 remkey = (key, i)
488 if input_value < old_value:
489 # it wrapped!
490 self.reminders[name][remkey] += old_value
491 self.reminder_keys[name][key].add(remkey)
492 bits.append(input_value + self.reminders[name][remkey])
493
494 new_dict[key] = tuple(bits)
495
496 self.cache[name] = input_dict
497 return new_dict
498
499 def cache_clear(self, name=None):
500 """Clear the internal cache, optionally only for function 'name'."""
501 with self.lock:
502 if name is None:
503 self.cache.clear()
504 self.reminders.clear()
505 self.reminder_keys.clear()
506 else:
507 self.cache.pop(name, None)
508 self.reminders.pop(name, None)
509 self.reminder_keys.pop(name, None)
510
511 def cache_info(self):
512 """Return internal cache dicts as a tuple of 3 elements."""
513 with self.lock:
514 return (self.cache, self.reminders, self.reminder_keys)
515
516
517def wrap_numbers(input_dict, name):
518 """Given an `input_dict` and a function `name`, adjust the numbers
519 which "wrap" (restart from zero) across different calls by adding
520 "old value" to "new value" and return an updated dict.
521 """
522 with _wn.lock:
523 return _wn.run(input_dict, name)
524
525
526_wn = _WrapNumbers()
527wrap_numbers.cache_clear = _wn.cache_clear
528wrap_numbers.cache_info = _wn.cache_info
529
530
531# The read buffer size for open() builtin. This (also) dictates how
532# much data we read(2) when iterating over file lines as in:
533# >>> with open(file) as f:
534# ... for line in f:
535# ... ...
536# Default per-line buffer size for binary files is 1K. For text files
537# is 8K. We use a bigger buffer (32K) in order to have more consistent
538# results when reading /proc pseudo files on Linux, see:
539# https://github.com/giampaolo/psutil/issues/2050
540# https://github.com/giampaolo/psutil/issues/708
541FILE_READ_BUFFER_SIZE = 32 * 1024
542
543
544def open_binary(fname):
545 return open(fname, "rb", buffering=FILE_READ_BUFFER_SIZE)
546
547
548def open_text(fname):
549 """Open a file in text mode by using the proper FS encoding and
550 en/decoding error handlers.
551 """
552 # See:
553 # https://github.com/giampaolo/psutil/issues/675
554 # https://github.com/giampaolo/psutil/pull/733
555 fobj = open( # noqa: SIM115
556 fname,
557 buffering=FILE_READ_BUFFER_SIZE,
558 encoding=ENCODING,
559 errors=ENCODING_ERRS,
560 )
561 try:
562 # Dictates per-line read(2) buffer size. Defaults is 8k. See:
563 # https://github.com/giampaolo/psutil/issues/2050#issuecomment-1013387546
564 fobj._CHUNK_SIZE = FILE_READ_BUFFER_SIZE
565 except AttributeError:
566 pass
567 except Exception:
568 fobj.close()
569 raise
570
571 return fobj
572
573
574def cat(fname, fallback=_DEFAULT, _open=open_text):
575 """Read entire file content and return it as a string. File is
576 opened in text mode. If specified, `fallback` is the value
577 returned in case of error, either if the file does not exist or
578 it can't be read().
579 """
580 if fallback is _DEFAULT:
581 with _open(fname) as f:
582 return f.read()
583 else:
584 try:
585 with _open(fname) as f:
586 return f.read()
587 except OSError:
588 return fallback
589
590
591def bcat(fname, fallback=_DEFAULT):
592 """Same as above but opens file in binary mode."""
593 return cat(fname, fallback=fallback, _open=open_binary)
594
595
596def bytes2human(n: int) -> str:
597 """Convert n bytes to a human-readable string.
598
599 >>> bytes2human(10000)
600 '9.8K'
601 >>> bytes2human(100001221)
602 '95.4M'
603 """
604 symbols = ('B', 'K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
605 if abs(n) < 1024:
606 return f"{int(n)}{symbols[0]}"
607 value = float(n)
608 for symbol in symbols[1:-1]:
609 value /= 1024
610 # 1023.96K would print as 1024.0K
611 if abs(round(value, 1)) < 1024:
612 return f"{value:.1f}{symbol}"
613 return f"{value / 1024:.1f}{symbols[-1]}"
614
615
616def decode(s):
617 return s.decode(encoding=ENCODING, errors=ENCODING_ERRS)
618
619
620# =====================================================================
621# --- shell utils
622# =====================================================================
623
624
625@functools.lru_cache
626def term_supports_colors(force_color=False):
627 if WINDOWS:
628 return False
629 if force_color:
630 return True
631 if not hasattr(sys.stdout, "isatty") or not sys.stdout.isatty():
632 return False
633 try:
634 sys.stdout.fileno()
635 except Exception: # noqa: BLE001
636 return False
637 return True
638
639
640def hilite(s, color=None, bold=False, force_color=False):
641 """Return an highlighted version of 'string'."""
642 if not term_supports_colors(force_color=force_color):
643 return s
644 attr = []
645 colors = dict(
646 blue='34',
647 brown='33',
648 darkgrey='30',
649 green='32',
650 grey='37',
651 lightblue='36',
652 red='31',
653 violet='35',
654 yellow='93',
655 )
656 colors[None] = '29'
657 try:
658 color = colors[color]
659 except KeyError:
660 msg = f"invalid color {color!r}; choose amongst {list(colors)}"
661 raise ValueError(msg) from None
662 attr.append(color)
663 if bold:
664 attr.append('1')
665 return f"\x1b[{';'.join(attr)}m{s}\x1b[0m"
666
667
668def print_color(
669 s, color=None, bold=False, file=sys.stdout
670): # pragma: no cover
671 """Print a colorized version of string."""
672 if term_supports_colors():
673 s = hilite(s, color=color, bold=bold)
674 print(s, file=file, flush=True)
675
676
677def debug(msg):
678 """If PSUTIL_DEBUG env var is set, print a debug message to stderr."""
679 if PSUTIL_DEBUG:
680 import inspect
681
682 fname, lineno, _, _lines, _index = inspect.getframeinfo(
683 inspect.currentframe().f_back
684 )
685 if isinstance(msg, Exception):
686 if isinstance(msg, OSError):
687 # ...because str(exc) may contain info about the file name
688 msg = f"ignoring {msg}"
689 else:
690 msg = f"ignoring {msg!r}"
691 print( # noqa: T201
692 f"psutil-debug [{fname}:{lineno}]> {msg}", file=sys.stderr
693 )
694
695
696def warn(msg):
697 """Emit a RuntimeWarning. Use it for events which are never
698 supposed to happen, and imply a psutil bug.
699 """
700 import inspect
701
702 fname, lineno, _, _lines, _index = inspect.getframeinfo(
703 inspect.currentframe().f_back
704 )
705 msg = f"{msg} (originated from {fname}:{lineno})"
706 if PSUTIL_TESTING:
707 msg = f"CRITICAL: {msg}"
708 raise RuntimeError(msg)
709 try:
710 warnings.warn(msg, RuntimeWarning, stacklevel=2)
711 except RuntimeWarning:
712 pass # -W error: we never want to fail because of a warning