Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
321 lines (263 sloc) 12.2 KB
#!/usr/bin/env python
# coding: utf-8
''' This is part of the MSS Python's module.
Source: https://github.com/BoboTiG/python-mss
'''
from ctypes import (
POINTER, Structure, byref, c_char_p, c_int, c_int32, c_long, c_uint,
c_uint32, c_ulong, c_ushort, c_void_p, cast, cdll, create_string_buffer)
from ctypes.util import find_library
from os import environ
from os.path import abspath, dirname, isfile, realpath
from struct import pack
from sys import maxsize, version
from .base import MSSBase
from .exception import ScreenshotError
__all__ = ['MSS']
class Display(Structure):
''' Structure that serves as the connection to the X server
and that contains all the information about that X server.
'''
class XWindowAttributes(Structure):
''' Attributes for the specified window. '''
_fields_ = [('x', c_int32), ('y', c_int32), ('width', c_int32),
('height', c_int32), ('border_width', c_int32),
('depth', c_int32), ('visual', c_ulong), ('root', c_ulong),
('class', c_int32), ('bit_gravity', c_int32),
('win_gravity', c_int32), ('backing_store', c_int32),
('backing_planes', c_ulong), ('backing_pixel', c_ulong),
('save_under', c_int32), ('colourmap', c_ulong),
('mapinstalled', c_uint32), ('map_state', c_uint32),
('all_event_masks', c_ulong), ('your_event_mask', c_ulong),
('do_not_propagate_mask', c_ulong),
('override_redirect', c_int32), ('screen', c_ulong)]
class XImage(Structure):
''' Description of an image as it exists in the client's memory.
https://tronche.com/gui/x/xlib/graphics/images.html
'''
_fields_ = [('width', c_int), ('height', c_int), ('xoffset', c_int),
('format', c_int), ('data', c_void_p),
('byte_order', c_int), ('bitmap_unit', c_int),
('bitmap_bit_order', c_int), ('bitmap_pad', c_int),
('depth', c_int), ('bytes_per_line', c_int),
('bits_per_pixel', c_int), ('red_mask', c_ulong),
('green_mask', c_ulong), ('blue_mask', c_ulong)]
class XRRModeInfo(Structure):
''' Voilà, voilà. '''
class XRRScreenResources(Structure):
''' Structure that contains arrays of XIDs that point to the
available outputs and associated CRTCs.
'''
_fields_ = [('timestamp', c_ulong), ('configTimestamp', c_ulong),
('ncrtc', c_int), ('crtcs', POINTER(c_long)),
('noutput', c_int), ('outputs', POINTER(c_long)),
('nmode', c_int), ('modes', POINTER(XRRModeInfo))]
class XRRCrtcInfo(Structure):
''' Structure that contains CRTC informations. '''
_fields_ = [('timestamp', c_ulong), ('x', c_int), ('y', c_int),
('width', c_int), ('height', c_int), ('mode', c_long),
('rotation', c_int), ('noutput', c_int),
('outputs', POINTER(c_long)), ('rotations', c_ushort),
('npossible', c_int), ('possible', POINTER(c_long))]
class MSS(MSSBase):
''' Mutliple ScreenShots implementation for GNU/Linux.
It uses intensively the Xlib and Xrandr extension.
'''
# pylint: disable=too-many-instance-attributes
display = None
use_mss = False
mss = None
xlib = None
xrandr = None
display = None
root = None
def __del__(self):
''' Disconnect from X server. '''
if self.display:
self.xlib.XCloseDisplay(self.display)
self.display = None
def __init__(self, display=None):
''' GNU/Linux initialisations. '''
if not display:
try:
if version > '3':
display = bytes(environ['DISPLAY'], 'utf-8')
else:
display = environ['DISPLAY']
except KeyError:
err = '$DISPLAY not set. Stopping to prevent segfault.'
raise ScreenshotError(err)
x11 = find_library('X11')
if not x11:
raise ScreenshotError('No X11 library found.')
self.xlib = cdll.LoadLibrary(x11)
xrandr = find_library('Xrandr')
if not xrandr:
raise ScreenshotError('No Xrandr extension found.')
self.xrandr = cdll.LoadLibrary(xrandr)
# libmss = find_library('mss')
libmss = '{0}/linux/{1}/libmss.so'.format(
dirname(realpath(abspath(__file__))), arch())
if isfile(libmss):
self.mss = cdll.LoadLibrary(libmss)
self.use_mss = True
else:
print('No MSS library found. Using slow native function.')
self._set_argtypes()
self._set_restypes()
self.display = self.xlib.XOpenDisplay(display)
try:
assert self.display.contents
except ValueError:
raise ScreenshotError('Cannot open display "{0}".'.format(display))
screen = self.xlib.XDefaultScreen(self.display)
self.root = self.xlib.XDefaultRootWindow(self.display, screen)
def _set_argtypes(self):
''' Functions arguments.
Curiously, if we set up XGetPixel arguments type,
the entire process takes twice more time.
So, no need to waste this precious time :)
Note: this issue does not occur when using libmss.
'''
self.xlib.XOpenDisplay.argtypes = [c_char_p]
self.xlib.XDefaultScreen.argtypes = [POINTER(Display)]
self.xlib.XDefaultRootWindow.argtypes = [POINTER(Display), c_int]
self.xlib.XGetWindowAttributes.argtypes = [POINTER(Display),
POINTER(XWindowAttributes),
POINTER(XWindowAttributes)]
self.xlib.XAllPlanes.argtypes = []
self.xlib.XGetImage.argtypes = [POINTER(Display), POINTER(Display),
c_int, c_int, c_uint, c_uint, c_ulong,
c_int]
# self.xlib.XGetPixel.argtypes = [POINTER(XImage), c_int, c_int]
self.xlib.XDestroyImage.argtypes = [POINTER(XImage)]
self.xlib.XCloseDisplay.argtypes = [POINTER(Display)]
self.xrandr.XRRGetScreenResources.argtypes = [POINTER(Display),
POINTER(Display)]
self.xrandr.XRRGetCrtcInfo.argtypes = [POINTER(Display),
POINTER(XRRScreenResources),
c_long]
self.xrandr.XRRFreeScreenResources.argtypes = \
[POINTER(XRRScreenResources)]
self.xrandr.XRRFreeCrtcInfo.argtypes = [POINTER(XRRCrtcInfo)]
if self.use_mss:
self.mss.GetXImagePixels.argtypes = [POINTER(XImage), c_void_p]
def _set_restypes(self):
''' Functions return type. '''
def validate(value):
''' Validate the returned value of xrandr.XRRGetScreenResources().
We can end on a segfault if not:
Xlib: extension "RANDR" missing on display "...".
'''
if value == 0:
err = 'xrandr.XRRGetScreenResources() failed.'
err += ' NULL pointer received.'
raise ScreenshotError(err)
return cast(value, POINTER(XRRScreenResources))
self.xlib.XOpenDisplay.restype = POINTER(Display)
self.xlib.XDefaultScreen.restype = c_int
self.xlib.XGetWindowAttributes.restype = c_int
self.xlib.XAllPlanes.restype = c_ulong
self.xlib.XGetImage.restype = POINTER(XImage)
self.xlib.XGetPixel.restype = c_ulong
self.xlib.XDestroyImage.restype = c_void_p
self.xlib.XCloseDisplay.restype = c_void_p
self.xlib.XDefaultRootWindow.restype = POINTER(XWindowAttributes)
self.xrandr.XRRGetScreenResources.restype = validate
self.xrandr.XRRGetCrtcInfo.restype = POINTER(XRRCrtcInfo)
self.xrandr.XRRFreeScreenResources.restype = c_void_p
self.xrandr.XRRFreeCrtcInfo.restype = c_void_p
if self.use_mss:
self.mss.GetXImagePixels.restype = c_int
def enum_display_monitors(self, force=False):
''' Get positions of monitors (see parent class). '''
if not self.monitors or force:
self.monitors = []
# All monitors
gwa = XWindowAttributes()
self.xlib.XGetWindowAttributes(self.display, self.root, byref(gwa))
self.monitors.append({
'left': int(gwa.x),
'top': int(gwa.y),
'width': int(gwa.width),
'height': int(gwa.height),
'monitor': 0
})
# Each monitors
# Fix for XRRGetScreenResources:
# expected LP_Display instance instead of LP_XWindowAttributes
root = cast(self.root, POINTER(Display))
mon = self.xrandr.XRRGetScreenResources(self.display, root)
for num in range(mon.contents.ncrtc):
crtc = self.xrandr.XRRGetCrtcInfo(self.display, mon,
mon.contents.crtcs[num])
self.monitors.append({
'left': int(crtc.contents.x),
'top': int(crtc.contents.y),
'width': int(crtc.contents.width),
'height': int(crtc.contents.height),
'monitor': num
})
self.xrandr.XRRFreeCrtcInfo(crtc)
self.xrandr.XRRFreeScreenResources(mon)
return self.monitors
def get_pixels(self, monitor):
''' Retrieve all pixels from a monitor. Pixels have to be RGB. '''
self.width = monitor['width']
self.height = monitor['height']
left, top = monitor['left'], monitor['top']
zpixmap = 2
allplanes = self.xlib.XAllPlanes()
# Fix for XGetImage:
# expected LP_Display instance instead of LP_XWindowAttributes
root = cast(self.root, POINTER(Display))
ximage = self.xlib.XGetImage(self.display, root, left, top,
self.width, self.height, allplanes,
zpixmap)
if not ximage:
err = 'xlib.XGetImage() failed. Monitor informations: '
for key, val in sorted(monitor.items()):
err = '{0}{1}: {2}, '.format(err, key, val)
err = err.strip(', ')
raise ScreenshotError(err)
if not self.use_mss:
self.get_pixels_slow(ximage)
else:
self.image = create_string_buffer(self.height * self.width * 3)
ret = self.mss.GetXImagePixels(ximage, self.image)
if not ret:
self.xlib.XDestroyImage(ximage)
err = 'libmss.GetXImagePixels() failed (retcode={0}).'
raise ScreenshotError(err.format(ret))
self.xlib.XDestroyImage(ximage)
return self.image
def get_pixels_slow(self, ximage):
''' Retrieve all pixels from a monitor. Pixels have to be RGB.
(!) Insanely slow version, see doc/linux-slow-version. (!)
'''
# @TODO: this part takes most of the time. Need a better solution.
def pix(pixel, _resultats={}, p__=pack):
''' Apply shifts to a pixel to get the RGB values.
This method uses of memoization.
'''
# pylint: disable=dangerous-default-value
if pixel not in _resultats:
_resultats[pixel] = p__('<B', (pixel & rmask) >> 16) + \
p__('<B', (pixel & gmask) >> 8) + \
p__('<B', pixel & bmask)
return _resultats[pixel]
self.width = ximage.contents.width
self.height = ximage.contents.height
rmask = ximage.contents.red_mask
bmask = ximage.contents.blue_mask
gmask = ximage.contents.green_mask
get_pix = self.xlib.XGetPixel
pixels = [pix(get_pix(ximage, x, y))
for y in range(self.height) for x in range(self.width)]
self.image = b''.join(pixels)
return self.image
def arch():
''' Detect OS architecture.
Returns an int: 32 or 64
'''
return 64 if maxsize > 2 ** 32 else 32
You can’t perform that action at this time.