"""
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.")