| """ |
| A wrapper around the Direct Rendering Manager (DRM) library, which itself is a |
| wrapper around the Direct Rendering Interface (DRI) between the kernel and |
| userland. |
| |
| Since we are masochists, we use ctypes instead of cffi to load libdrm and |
| access several symbols within it. We use Python's file descriptor and mmap |
| wrappers. |
| |
| At some point in the future, cffi could be used, for approximately the same |
| cost in lines of code. |
| """ |
| |
| from ctypes import * |
| import exceptions |
| import mmap |
| import os |
| import subprocess |
| |
| from PIL import Image |
| |
| # drmModeConnection enum |
| DRM_MODE_CONNECTED = 1 |
| DRM_MODE_DISCONNECTED = 2 |
| DRM_MODE_UNKNOWNCONNECTION = 3 |
| |
| DRM_MODE_CONNECTOR_Unknown = 0 |
| DRM_MODE_CONNECTOR_VGA = 1 |
| DRM_MODE_CONNECTOR_DVII = 2 |
| DRM_MODE_CONNECTOR_DVID = 3 |
| DRM_MODE_CONNECTOR_DVIA = 4 |
| DRM_MODE_CONNECTOR_Composite = 5 |
| DRM_MODE_CONNECTOR_SVIDEO = 6 |
| DRM_MODE_CONNECTOR_LVDS = 7 |
| DRM_MODE_CONNECTOR_Component = 8 |
| DRM_MODE_CONNECTOR_9PinDIN = 9 |
| DRM_MODE_CONNECTOR_DisplayPort = 10 |
| DRM_MODE_CONNECTOR_HDMIA = 11 |
| DRM_MODE_CONNECTOR_HDMIB = 12 |
| DRM_MODE_CONNECTOR_TV = 13 |
| DRM_MODE_CONNECTOR_eDP = 14 |
| DRM_MODE_CONNECTOR_VIRTUAL = 15 |
| DRM_MODE_CONNECTOR_DSI = 16 |
| |
| |
| class DrmVersion(Structure): |
| """ |
| The version of a DRM node. |
| """ |
| |
| _fields_ = [ |
| ("version_major", c_int), |
| ("version_minor", c_int), |
| ("version_patchlevel", c_int), |
| ("name_len", c_int), |
| ("name", c_char_p), |
| ("date_len", c_int), |
| ("date", c_char_p), |
| ("desc_len", c_int), |
| ("desc", c_char_p), |
| ] |
| |
| _l = None |
| |
| def __repr__(self): |
| return "%s %d.%d.%d (%s) (%s)" % (self.name, |
| self.version_major, |
| self.version_minor, |
| self.version_patchlevel, |
| self.desc, |
| self.date,) |
| |
| def __del__(self): |
| if self._l: |
| self._l.drmFreeVersion(self) |
| |
| |
| class DrmModeResources(Structure): |
| """ |
| Resources associated with setting modes on a DRM node. |
| """ |
| |
| _fields_ = [ |
| ("count_fbs", c_int), |
| ("fbs", POINTER(c_uint)), |
| ("count_crtcs", c_int), |
| ("crtcs", POINTER(c_uint)), |
| ("count_connectors", c_int), |
| ("connectors", POINTER(c_uint)), |
| ("count_encoders", c_int), |
| ("encoders", POINTER(c_uint)), |
| ("min_width", c_int), |
| ("max_width", c_int), |
| ("min_height", c_int), |
| ("max_height", c_int), |
| ] |
| |
| _fd = None |
| _l = None |
| |
| def __repr__(self): |
| return "<DRM mode resources>" |
| |
| def __del__(self): |
| if self._l: |
| self._l.drmModeFreeResources(self) |
| |
| def _wakeup_screen(self): |
| """ |
| Send a synchronous dbus message to power on screen. |
| """ |
| # Get and process reply to make this synchronous. |
| subprocess.check_output([ |
| "dbus-send", "--type=method_call", "--system", "--print-reply", |
| "--dest=org.chromium.PowerManager", "/org/chromium/PowerManager", |
| "org.chromium.PowerManager.HandleUserActivity", "int32:0" |
| ]) |
| |
| def getValidCrtc(self): |
| for i in xrange(0, self.count_crtcs): |
| crtc_id = self.crtcs[i] |
| crtc = self._l.drmModeGetCrtc(self._fd, crtc_id).contents |
| if crtc.mode_valid: |
| return crtc |
| return None |
| |
| def getCrtc(self, crtc_id): |
| """ |
| Obtain the CRTC at a given index. |
| |
| @param crtc_id: The CRTC to get. |
| """ |
| if crtc_id: |
| return self._l.drmModeGetCrtc(self._fd, crtc_id).contents |
| return self.getValidCrtc() |
| |
| def getCrtcRobust(self, crtc_id=None): |
| crtc = self.getCrtc(crtc_id) |
| if crtc is None: |
| self._wakeup_screen() |
| crtc = self.getCrtc(crtc_id) |
| if crtc is not None: |
| crtc._fd = self._fd |
| crtc._l = self._l |
| return crtc |
| |
| |
| class DrmModeModeInfo(Structure): |
| """ |
| A DRM modesetting mode info. |
| """ |
| |
| _fields_ = [ |
| ("clock", c_uint), |
| ("hdisplay", c_ushort), |
| ("hsync_start", c_ushort), |
| ("hsync_end", c_ushort), |
| ("htotal", c_ushort), |
| ("hskew", c_ushort), |
| ("vdisplay", c_ushort), |
| ("vsync_start", c_ushort), |
| ("vsync_end", c_ushort), |
| ("vtotal", c_ushort), |
| ("vscan", c_ushort), |
| ("vrefresh", c_uint), |
| ("flags", c_uint), |
| ("type", c_uint), |
| ("name", c_char * 32), |
| ] |
| |
| |
| class DrmModeCrtc(Structure): |
| """ |
| A DRM modesetting CRTC. |
| """ |
| |
| _fields_ = [ |
| ("crtc_id", c_uint), |
| ("buffer_id", c_uint), |
| ("x", c_uint), |
| ("y", c_uint), |
| ("width", c_uint), |
| ("height", c_uint), |
| ("mode_valid", c_int), |
| ("mode", DrmModeModeInfo), |
| ("gamma_size", c_int), |
| ] |
| |
| _fd = None |
| _l = None |
| |
| def __repr__(self): |
| return "<CRTC (%d)>" % self.crtc_id |
| |
| def __del__(self): |
| if self._l: |
| self._l.drmModeFreeCrtc(self) |
| |
| def hasFb(self): |
| """ |
| Whether this CRTC has an associated framebuffer. |
| """ |
| |
| return self.buffer_id != 0 |
| |
| def fb(self): |
| """ |
| Obtain the framebuffer, if one is associated. |
| """ |
| |
| if self.hasFb(): |
| fb = self._l.drmModeGetFB(self._fd, self.buffer_id).contents |
| fb._fd = self._fd |
| fb._l = self._l |
| return fb |
| else: |
| raise RuntimeError("CRTC %d doesn't have a framebuffer!" % |
| self.crtc_id) |
| |
| |
| class DrmModeEncoder(Structure): |
| """ |
| A DRM modesetting encoder. |
| """ |
| |
| _fields_ = [ |
| ("encoder_id", c_uint), |
| ("encoder_type", c_uint), |
| ("crtc_id", c_uint), |
| ("possible_crtcs", c_uint), |
| ("possible_clones", c_uint), |
| ] |
| |
| _fd = None |
| _l = None |
| |
| def __repr__(self): |
| return "<Encoder (%d)>" % self.encoder_id |
| |
| def __del__(self): |
| if self._l: |
| self._l.drmModeFreeEncoder(self) |
| |
| |
| class DrmModeConnector(Structure): |
| """ |
| A DRM modesetting connector. |
| """ |
| |
| _fields_ = [ |
| ("connector_id", c_uint), |
| ("encoder_id", c_uint), |
| ("connector_type", c_uint), |
| ("connector_type_id", c_uint), |
| ("connection", c_uint), # drmModeConnection enum |
| ("mmWidth", c_uint), |
| ("mmHeight", c_uint), |
| ("subpixel", c_uint), # drmModeSubPixel enum |
| ("count_modes", c_int), |
| ("modes", POINTER(DrmModeModeInfo)), |
| ("count_propts", c_int), |
| ("props", POINTER(c_uint)), |
| ("prop_values", POINTER(c_ulonglong)), |
| ("count_encoders", c_int), |
| ("encoders", POINTER(c_uint)), |
| ] |
| |
| _fd = None |
| _l = None |
| |
| def __repr__(self): |
| return "<Connector (%d)>" % self.connector_id |
| |
| def __del__(self): |
| if self._l: |
| self._l.drmModeFreeConnector(self) |
| |
| def isInternal(self): |
| return (self.connector_type == DRM_MODE_CONNECTOR_LVDS or |
| self.connector_type == DRM_MODE_CONNECTOR_eDP or |
| self.connector_type == DRM_MODE_CONNECTOR_DSI) |
| |
| def isConnected(self): |
| return self.connection == DRM_MODE_CONNECTED |
| |
| |
| class drm_mode_map_dumb(Structure): |
| """ |
| Request a mapping of a modesetting buffer. |
| |
| The map will be "dumb;" it will be accessible via mmap() but very slow. |
| """ |
| |
| _fields_ = [ |
| ("handle", c_uint), |
| ("pad", c_uint), |
| ("offset", c_ulonglong), |
| ] |
| |
| |
| # This constant is not defined in any one header; it is the pieced-together |
| # incantation for the ioctl that performs dumb mappings. I would love for this |
| # to not have to be here, but it can't be imported from any header easily. |
| DRM_IOCTL_MODE_MAP_DUMB = 0xc01064b3 |
| |
| |
| class DrmModeFB(Structure): |
| """ |
| A DRM modesetting framebuffer. |
| """ |
| |
| _fields_ = [ |
| ("fb_id", c_uint), |
| ("width", c_uint), |
| ("height", c_uint), |
| ("pitch", c_uint), |
| ("bpp", c_uint), |
| ("depth", c_uint), |
| ("handle", c_uint), |
| ] |
| |
| _l = None |
| _map = None |
| |
| def __repr__(self): |
| s = "<Framebuffer (%dx%d (pitch %d bytes), %d bits/pixel, depth %d)" |
| vitals = s % (self.width, |
| self.height, |
| self.pitch, |
| self.bpp, |
| self.depth,) |
| if self._map: |
| tail = " (mapped)>" |
| else: |
| tail = ">" |
| return vitals + tail |
| |
| def __del__(self): |
| if self._l: |
| self._l.drmModeFreeFB(self) |
| |
| def map(self, size): |
| """ |
| Map the framebuffer. |
| """ |
| |
| if self._map: |
| return |
| |
| mapDumb = drm_mode_map_dumb() |
| mapDumb.handle = self.handle |
| |
| rv = self._l.drmIoctl(self._fd, DRM_IOCTL_MODE_MAP_DUMB, |
| pointer(mapDumb)) |
| if rv: |
| raise IOError(rv, os.strerror(rv)) |
| |
| # mmap.mmap() has a totally different order of arguments in Python |
| # compared to C; check the documentation before altering this |
| # incantation. |
| self._map = mmap.mmap(self._fd, |
| size, |
| flags=mmap.MAP_SHARED, |
| prot=mmap.PROT_READ, |
| offset=mapDumb.offset) |
| |
| def unmap(self): |
| """ |
| Unmap the framebuffer. |
| """ |
| |
| if self._map: |
| self._map.close() |
| self._map = None |
| |
| |
| def loadDRM(): |
| """ |
| Load a handle to libdrm. |
| |
| In addition to loading, this function also configures the argument and |
| return types of functions. |
| """ |
| |
| l = None |
| |
| try: |
| l = cdll.LoadLibrary("libdrm.so") |
| except OSError: |
| l = cdll.LoadLibrary("libdrm.so.2") # ubuntu doesn't have libdrm.so |
| |
| l.drmGetVersion.argtypes = [c_int] |
| l.drmGetVersion.restype = POINTER(DrmVersion) |
| |
| l.drmFreeVersion.argtypes = [POINTER(DrmVersion)] |
| l.drmFreeVersion.restype = None |
| |
| l.drmModeGetResources.argtypes = [c_int] |
| l.drmModeGetResources.restype = POINTER(DrmModeResources) |
| |
| l.drmModeFreeResources.argtypes = [POINTER(DrmModeResources)] |
| l.drmModeFreeResources.restype = None |
| |
| l.drmModeGetCrtc.argtypes = [c_int, c_uint] |
| l.drmModeGetCrtc.restype = POINTER(DrmModeCrtc) |
| |
| l.drmModeFreeCrtc.argtypes = [POINTER(DrmModeCrtc)] |
| l.drmModeFreeCrtc.restype = None |
| |
| l.drmModeGetEncoder.argtypes = [c_int, c_uint] |
| l.drmModeGetEncoder.restype = POINTER(DrmModeEncoder) |
| |
| l.drmModeFreeEncoder.argtypes = [POINTER(DrmModeEncoder)] |
| l.drmModeFreeEncoder.restype = None |
| |
| l.drmModeGetConnector.argtypes = [c_int, c_uint] |
| l.drmModeGetConnector.restype = POINTER(DrmModeConnector) |
| |
| l.drmModeFreeConnector.argtypes = [POINTER(DrmModeConnector)] |
| l.drmModeFreeConnector.restype = None |
| |
| l.drmModeGetFB.argtypes = [c_int, c_uint] |
| l.drmModeGetFB.restype = POINTER(DrmModeFB) |
| |
| l.drmModeFreeFB.argtypes = [POINTER(DrmModeFB)] |
| l.drmModeFreeFB.restype = None |
| |
| l.drmIoctl.argtypes = [c_int, c_ulong, c_voidp] |
| l.drmIoctl.restype = c_int |
| |
| return l |
| |
| |
| class DRM(object): |
| """ |
| A DRM node. |
| """ |
| |
| def __init__(self, library, fd): |
| self._l = library |
| self._fd = fd |
| |
| def __repr__(self): |
| return "<DRM (FD %d)>" % self._fd |
| |
| @classmethod |
| def fromHandle(cls, handle): |
| """ |
| Create a node from a file handle. |
| |
| @param handle: A file-like object backed by a file descriptor. |
| """ |
| |
| self = cls(loadDRM(), handle.fileno()) |
| # We must keep the handle alive, and we cannot trust the caller to |
| # keep it alive for us. |
| self._handle = handle |
| return self |
| |
| def version(self): |
| """ |
| Obtain the version. |
| """ |
| |
| v = self._l.drmGetVersion(self._fd).contents |
| v._l = self._l |
| return v |
| |
| def resources(self): |
| """ |
| Obtain the modesetting resources. |
| """ |
| |
| resources_ptr = self._l.drmModeGetResources(self._fd) |
| if resources_ptr: |
| r = resources_ptr.contents |
| r._fd = self._fd |
| r._l = self._l |
| return r |
| |
| return None |
| |
| def getCrtc(self, crtc_id): |
| c_ptr = self._l.drmModeGetCrtc(self._fd, crtc_id) |
| if c_ptr: |
| c = c_ptr.contents |
| c._fd = self._fd |
| c._l = self._l |
| return c |
| |
| return None |
| |
| def getEncoder(self, encoder_id): |
| e_ptr = self._l.drmModeGetEncoder(self._fd, encoder_id) |
| if e_ptr: |
| e = e_ptr.contents |
| e._fd = self._fd |
| e._l = self._l |
| return e |
| |
| return None |
| |
| def getConnector(self, connector_id): |
| c_ptr = self._l.drmModeGetConnector(self._fd, connector_id) |
| if c_ptr: |
| c = c_ptr.contents |
| c._fd = self._fd |
| c._l = self._l |
| return c |
| |
| return None |
| |
| |
| |
| def drmFromPath(path): |
| """ |
| Given a DRM node path, open the corresponding node. |
| |
| @param path: The path of the minor node to open. |
| """ |
| |
| handle = open(path) |
| return DRM.fromHandle(handle) |
| |
| |
| def _bgrx24(i): |
| b = ord(next(i)) |
| g = ord(next(i)) |
| r = ord(next(i)) |
| next(i) |
| return r, g, b |
| |
| |
| def _copyImageBlocklinear(image, fb, unformat): |
| gobPitch = 64 |
| gobHeight = 128 |
| while gobHeight > 8 and gobHeight >= 2 * fb.height: |
| gobHeight //= 2 |
| gobSize = gobPitch * gobHeight |
| gobWidth = gobPitch // (fb.bpp // 8) |
| |
| gobCountX = (fb.pitch + gobPitch - 1) // gobPitch |
| gobCountY = (fb.height + gobHeight - 1) // gobHeight |
| fb.map(gobCountX * gobCountY * gobSize) |
| m = fb._map |
| |
| offset = 0 |
| for gobY in range(gobCountY): |
| gobTop = gobY * gobHeight |
| for gobX in range(gobCountX): |
| m.seek(offset) |
| gob = m.read(gobSize) |
| iterGob = iter(gob) |
| gobLeft = gobX * gobWidth |
| for i in range(gobWidth * gobHeight): |
| rgb = unformat(iterGob) |
| x = gobLeft + (((i >> 3) & 8) | ((i >> 1) & 4) | (i & 3)) |
| y = gobTop + ((i >> 7 << 3) | ((i >> 3) & 6) | ((i >> 2) & 1)) |
| if x < fb.width and y < fb.height: |
| image.putpixel((x, y), rgb) |
| offset += gobSize |
| fb.unmap() |
| |
| |
| def _copyImageLinear(image, fb, unformat): |
| fb.map(fb.pitch * fb.height) |
| m = fb._map |
| pitch = fb.pitch |
| lineLength = fb.width * fb.bpp // 8 |
| for y in range(fb.height): |
| offset = y * pitch |
| m.seek(offset) |
| channels = m.read(lineLength) |
| ichannels = iter(channels) |
| for x in range(fb.width): |
| rgb = unformat(ichannels) |
| image.putpixel((x, y), rgb) |
| fb.unmap() |
| |
| |
| def _screenshot(drm, image, fb): |
| if fb.depth == 24: |
| unformat = _bgrx24 |
| else: |
| raise RuntimeError("Couldn't unformat FB: %r" % fb) |
| |
| if drm.version().name == "tegra": |
| _copyImageBlocklinear(image, fb, unformat) |
| else: |
| _copyImageLinear(image, fb, unformat) |
| |
| |
| _drm = None |
| |
| |
| 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 |
| """ |
| |
| global _drm |
| |
| if not _drm: |
| paths = [ |
| "/dev/dri/" + n |
| for n in filter(lambda x: x.startswith("card"), |
| os.listdir("/dev/dri")) |
| ] |
| |
| if crtc_id == "usb" or crtc_id == "evdi" or crtc_id == "udl": |
| for p in paths: |
| d = drmFromPath(p) |
| v = d.version() |
| |
| if crtc_id == v.name: |
| _drm = d |
| break |
| |
| if crtc_id == "usb" and (v.name == "evdi" or v.name == "udl"): |
| _drm = d |
| break |
| |
| elif crtc_id == "internal" or crtc_id == "external": |
| internal = crtc_id == "internal" |
| for p in paths: |
| d = drmFromPath(p) |
| if d.resources() is None: |
| continue |
| if d.resources() and d.resources().count_connectors > 0: |
| for c in xrange(0, d.resources().count_connectors): |
| connector = d.getConnector(d.resources().connectors[c]) |
| if (internal == connector.isInternal() |
| and connector.isConnected() |
| and connector.encoder_id != 0): |
| e = d.getEncoder(connector.encoder_id) |
| crtc = d.getCrtc(e.crtc_id) |
| if crtc.mode_valid: |
| crtc_id = crtc.crtc_id |
| _drm = d |
| break |
| if _drm: |
| break |
| |
| elif crtc_id is None or crtc_id == 0: |
| for p in paths: |
| d = drmFromPath(p) |
| if d.resources() is None: |
| continue |
| for c in xrange(0, d.resources().count_crtcs): |
| crtc = d.getCrtc(d.resources().crtcs[c]) |
| if crtc.mode_valid: |
| crtc_id = d.resources().crtcs[c] |
| _drm = d |
| break |
| if _drm: |
| break |
| |
| else: |
| for p in paths: |
| d = drmFromPath(p) |
| if d.resources() is None: |
| continue |
| for c in xrange(0, d.resources().count_crtcs): |
| if crtc_id == d.resources().crtcs[c]: |
| _drm = d |
| break |
| if _drm: |
| break |
| |
| if _drm: |
| crtc = _drm.resources().getCrtcRobust(crtc_id) |
| if crtc is not None: |
| framebuffer = crtc.fb() |
| image = Image.new("RGB", (framebuffer.width, framebuffer.height)) |
| _screenshot(_drm, image, framebuffer) |
| return image |
| |
| raise RuntimeError( |
| "Unable to take screenshot. There may not be anything on the screen.") |
