"""A wrapper around the Generic Buffer Manager (GBM) library. Currently implements exactly the functions required to screenshot a frame buffer using DRM crtc info. """ from ctypes import * import drm from PIL import Image GBM_BO_IMPORT_FD = 0x5503 GBM_BO_USE_SCANOUT = c_uint(1) GBM_BO_TRANSFER_READ = c_uint(1) GBM_MAX_PLANES = 4 def __gbm_fourcc_code(a, b, c, d): return ord(a) | (ord(b) << 8) | (ord(c) << 16) | (ord(d) << 24) GBM_FORMAT_ARGB8888 = __gbm_fourcc_code("A", "R", "2", "4") GBM_LIBRARIES = ["libgbm.so", "libgbm.so.1"] class gbm_import_fd_data(Structure): _fields_ = [ ("fd", c_int), ("width", c_uint), ("height", c_uint), ("stride", c_uint), ("bo_format", c_uint), ] class gbm_import_fd_planar_data(Structure): _fields_ = [ ("fds", c_int * GBM_MAX_PLANES), ("width", c_uint), ("height", c_uint), ("bo_format", c_uint), ("strides", c_uint * GBM_MAX_PLANES), ("offsets", c_uint * GBM_MAX_PLANES), ("format_modifiers", c_ulonglong * GBM_MAX_PLANES), ] class gbm_device(Structure): """Opaque struct for GBM device. """ pass class gbm_bo(Structure): """Opaque struct for GBM buffer. """ pass def loadGBM(): """Load and return a handle to libgbm.so. """ l = None for lib in GBM_LIBRARIES: try: l = cdll.LoadLibrary(lib) except OSError: l = None if l is not None: break if l is None: raise RuntimeError("Could not load GBM library.") return None l.gbm_create_device.argtypes = [c_int] l.gbm_create_device.restype = POINTER(gbm_device) l.gbm_device_destroy.argtypes = [POINTER(gbm_device)] l.gbm_device_destroy.restype = None l.gbm_bo_import.argtypes = [POINTER(gbm_device), c_uint, c_void_p, c_uint] l.gbm_bo_import.restype = POINTER(gbm_bo) l.gbm_bo_map.argtypes = [ POINTER(gbm_bo), c_uint, c_uint, c_uint, c_uint, c_uint, POINTER(c_uint), POINTER(c_void_p), c_size_t ] l.gbm_bo_map.restype = c_void_p l.gbm_bo_unmap.argtypes = [POINTER(gbm_bo), c_void_p] l.gbm_bo_unmap.restype = None return l class GBMBuffer(object): """A GBM buffer. """ def __init__(self, library, buffer): self._l = library self._buffer = buffer @classmethod def fromFD(cls, device, fd, width, height, stride, bo_format, usage): """Create/import a GBM Buffer Object from a file descriptor. @param device: GBM device object. @param fd: a file descriptor for the buffer to be imported. @param width: buffer width in pixels. @param height: buffer height in pixels. @param stride: buffer pitch; number of pixels between sequential rows. @param bo_format: pixel format fourcc code, e.g. GBM_FORMAT_ARGB8888. @param usage: buffer usage flags, e.g. GBM_BO_USE_SCANOUT. """ bo_data = gbm_import_fd_data() bo_data.fd = fd bo_data.width = width bo_data.height = height bo_data.stride = stride bo_data.bo_format = bo_format buffer = device._l.gbm_bo_import(device._device, GBM_BO_IMPORT_FD, byref(bo_data), usage) if buffer is None: raise RuntimeError("gbm_bo_import() returned NULL") self = cls(device._l, buffer) return self def map(self, x, y, width, height, flags, plane): """Map buffer data into this user-space. Returns (address, stride_bytes): void* start address for pixel array, number of BYTES between sequental rows of pixels. @param flags: The union of the GBM_BO_TRANSFER_* flags for this buffer. """ self._map_p = c_void_p(0) stride_out = c_uint(0) if width == 0 or height == 0: raise RuntimeError("Map width and/or height is 0") map_p = self._l.gbm_bo_map(self._buffer, x, y, width, height, flags, byref(stride_out), byref(self._map_p), plane) if stride_out is 0: raise RuntimeError("gbm_bo_map() stride is 0") if map_p is 0: raise RuntimeError("gbm_bo_map() returned NULL") return map_p, stride_out def unmap(self, map_data_p): self._l.gbm_bo_unmap(self._buffer, map_data_p) class GBMDevice(object): """A GBM device. """ def __init__(self, library, handle): self._l = library self._handle = handle self._device = library.gbm_create_device(self._handle) def __del__(self): if self._l: self._l.gbm_device_destroy(self._device) @classmethod def fromHandle(cls, handle): """Create a device object from an open file descriptor. """ self = cls(loadGBM(), handle) return self def _bgrx24(i): b = i & 255 g = (i >> 8) & 255 r = (i >> 16) & 255 return r, g, b def _copyImage(image, map_ints, map_ints_pitch, unformat): width, height = image.size for y in range(height): y_offset = y * map_ints_pitch for x in range(width): rgb = unformat(map_ints[y_offset + x]) image.putpixel((x, y), rgb) def crtcScreenshot(crtc_id=None): """Take a screenshot, returning an image object. @param crtc_id: The CRTC to screenshot. None for first found CRTC with mode set or "internal" for crtc connected to internal LCD or "external" for crtc connected to external display or "usb" "evdi" or "udl" for crtc with valid mode on evdi or udl display or DRM integer crtc_id """ crtc = drm.getCrtc(crtc_id) if crtc is not None: device = GBMDevice.fromHandle(drm._drm._fd) framebuffer = crtc.fb() # TODO(djmk): The buffer format is hardcoded to ARGB8888, we should fix # this to query for the frambuffer's format instead. format_hardcode = GBM_FORMAT_ARGB8888 bo = GBMBuffer.fromFD(device, framebuffer.getFD(), framebuffer.width, framebuffer.height, framebuffer.pitch, format_hardcode, GBM_BO_USE_SCANOUT) map_void_p, stride_bytes = bo.map(0, 0, framebuffer.width, framebuffer.height, GBM_BO_TRANSFER_READ, 0) stride_pixels = stride_bytes.value / 4 map_ints_type = c_int * (stride_pixels * framebuffer.height) map_int_p = cast(map_void_p, POINTER(c_int)) addr = addressof(map_int_p.contents) map_ints = map_ints_type.from_address(addr) image = Image.new("RGB", (framebuffer.width, framebuffer.height)) _copyImage(image, map_ints, stride_pixels, _bgrx24) bo.unmap(bo._map_p) return image raise RuntimeError( "Unable to take screenshot. There may not be anything on the screen.")