• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1"""
2A wrapper around the Direct Rendering Manager (DRM) library, which itself is a
3wrapper around the Direct Rendering Interface (DRI) between the kernel and
4userland.
5
6Since we are masochists, we use ctypes instead of cffi to load libdrm and
7access several symbols within it. We use Python's file descriptor and mmap
8wrappers.
9
10At some point in the future, cffi could be used, for approximately the same
11cost in lines of code.
12"""
13
14from ctypes import *
15import mmap
16import os
17
18from PIL import Image
19
20
21class DrmVersion(Structure):
22    """
23    The version of a DRM node.
24    """
25
26    _fields_ = [
27        ("version_major", c_int),
28        ("version_minor", c_int),
29        ("version_patchlevel", c_int),
30        ("name_len", c_int),
31        ("name", c_char_p),
32        ("date_len", c_int),
33        ("date", c_char_p),
34        ("desc_len", c_int),
35        ("desc", c_char_p),
36    ]
37
38    _l = None
39
40    def __repr__(self):
41        return "%s %d.%d.%d (%s) (%s)" % (
42            self.name,
43            self.version_major,
44            self.version_minor,
45            self.version_patchlevel,
46            self.desc,
47            self.date,
48        )
49
50    def __del__(self):
51        if self._l:
52            self._l.drmFreeVersion(self)
53
54
55class DrmModeResources(Structure):
56    """
57    Resources associated with setting modes on a DRM node.
58    """
59
60    _fields_ = [
61        ("count_fbs", c_int),
62        ("fbs", POINTER(c_uint)),
63        ("count_crtcs", c_int),
64        ("crtcs", POINTER(c_uint)),
65        ("count_connectors", c_int),
66        ("connectors", POINTER(c_uint)),
67        ("count_encoders", c_int),
68        ("encoders", POINTER(c_uint)),
69        ("min_width", c_int),
70        ("max_width", c_int),
71        ("min_height", c_int),
72        ("max_height", c_int),
73    ]
74
75    _fd = None
76    _l = None
77
78    def __repr__(self):
79        return "<DRM mode resources>"
80
81    def __del__(self):
82        if self._l:
83            self._l.drmModeFreeResources(self)
84
85    def getValidCrtc(self):
86        for i in xrange(0, self.count_crtcs):
87            crtc_id = self.crtcs[i]
88            crtc = self._l.drmModeGetCrtc(self._fd, crtc_id).contents
89            if crtc.mode_valid:
90                return crtc
91        return None
92
93    def getCrtc(self, crtc_id=None):
94        """
95        Obtain the CRTC at a given index.
96
97        @param crtc_id: The CRTC to get.
98        """
99        crtc = None
100        if crtc_id:
101            crtc = self._l.drmModeGetCrtc(self._fd, crtc_id).contents
102        else:
103            crtc = self.getValidCrtc()
104        crtc._fd = self._fd
105        crtc._l = self._l
106        return crtc
107
108
109class DrmModeCrtc(Structure):
110    """
111    A DRM modesetting CRTC.
112    """
113
114    _fields_ = [
115        ("crtc_id", c_uint),
116        ("buffer_id", c_uint),
117        ("x", c_uint),
118        ("y", c_uint),
119        ("width", c_uint),
120        ("height", c_uint),
121        ("mode_valid", c_int),
122        # XXX incomplete struct!
123    ]
124
125    _fd = None
126    _l = None
127
128    def __repr__(self):
129        return "<CRTC (%d)>" % self.crtc_id
130
131    def __del__(self):
132        if self._l:
133            self._l.drmModeFreeCrtc(self)
134
135    def hasFb(self):
136        """
137        Whether this CRTC has an associated framebuffer.
138        """
139
140        return self.buffer_id != 0
141
142    def fb(self):
143        """
144        Obtain the framebuffer, if one is associated.
145        """
146
147        if self.hasFb():
148            fb = self._l.drmModeGetFB(self._fd, self.buffer_id).contents
149            fb._fd = self._fd
150            fb._l = self._l
151            return fb
152        else:
153            raise RuntimeError("CRTC %d doesn't have a framebuffer!" %
154                               self.crtc_id)
155
156
157class drm_mode_map_dumb(Structure):
158    """
159    Request a mapping of a modesetting buffer.
160
161    The map will be "dumb;" it will be accessible via mmap() but very slow.
162    """
163
164    _fields_ = [
165        ("handle", c_uint),
166        ("pad", c_uint),
167        ("offset", c_ulonglong),
168    ]
169
170
171# This constant is not defined in any one header; it is the pieced-together
172# incantation for the ioctl that performs dumb mappings. I would love for this
173# to not have to be here, but it can't be imported from any header easily.
174DRM_IOCTL_MODE_MAP_DUMB = 0xc01064b3
175
176
177class DrmModeFB(Structure):
178    """
179    A DRM modesetting framebuffer.
180    """
181
182    _fields_ = [
183        ("fb_id", c_uint),
184        ("width", c_uint),
185        ("height", c_uint),
186        ("pitch", c_uint),
187        ("bpp", c_uint),
188        ("depth", c_uint),
189        ("handle", c_uint),
190    ]
191
192    _l = None
193    _map = None
194
195    def __repr__(self):
196        s = "<Framebuffer (%dx%d (pitch %d bytes), %d bits/pixel, depth %d)"
197        vitals = s % (
198            self.width,
199            self.height,
200            self.pitch,
201            self.bpp,
202            self.depth,
203        )
204        if self._map:
205            tail = " (mapped)>"
206        else:
207            tail = ">"
208        return vitals + tail
209
210    def __del__(self):
211        if self._l:
212            self._l.drmModeFreeFB(self)
213
214    def map(self):
215        """
216        Map the framebuffer.
217        """
218
219        if self._map:
220            return
221
222        mapDumb = drm_mode_map_dumb()
223        mapDumb.handle = self.handle
224
225        rv = self._l.drmIoctl(self._fd, DRM_IOCTL_MODE_MAP_DUMB,
226                              pointer(mapDumb))
227        if rv:
228            raise IOError(rv, os.strerror(rv))
229
230        size = self.pitch * self.height
231
232        # mmap.mmap() has a totally different order of arguments in Python
233        # compared to C; check the documentation before altering this
234        # incantation.
235        self._map = mmap.mmap(self._fd, size, flags=mmap.MAP_SHARED,
236                              prot=mmap.PROT_READ, offset=mapDumb.offset)
237
238    def unmap(self):
239        """
240        Unmap the framebuffer.
241        """
242
243        if self._map:
244            self._map.close()
245            self._map = None
246
247
248def loadDRM():
249    """
250    Load a handle to libdrm.
251
252    In addition to loading, this function also configures the argument and
253    return types of functions.
254    """
255
256    l = cdll.LoadLibrary("libdrm.so")
257
258    l.drmGetVersion.argtypes = [c_int]
259    l.drmGetVersion.restype = POINTER(DrmVersion)
260
261    l.drmFreeVersion.argtypes = [POINTER(DrmVersion)]
262    l.drmFreeVersion.restype = None
263
264    l.drmModeGetResources.argtypes = [c_int]
265    l.drmModeGetResources.restype = POINTER(DrmModeResources)
266
267    l.drmModeFreeResources.argtypes = [POINTER(DrmModeResources)]
268    l.drmModeFreeResources.restype = None
269
270    l.drmModeGetCrtc.argtypes = [c_int, c_uint]
271    l.drmModeGetCrtc.restype = POINTER(DrmModeCrtc)
272
273    l.drmModeFreeCrtc.argtypes = [POINTER(DrmModeCrtc)]
274    l.drmModeFreeCrtc.restype = None
275
276    l.drmModeGetFB.argtypes = [c_int, c_uint]
277    l.drmModeGetFB.restype = POINTER(DrmModeFB)
278
279    l.drmModeFreeFB.argtypes = [POINTER(DrmModeFB)]
280    l.drmModeFreeFB.restype = None
281
282    l.drmIoctl.argtypes = [c_int, c_ulong, c_voidp]
283    l.drmIoctl.restype = c_int
284
285    return l
286
287
288class DRM(object):
289    """
290    A DRM node.
291    """
292
293    def __init__(self, library, fd):
294        self._l = library
295        self._fd = fd
296
297    def __repr__(self):
298        return "<DRM (FD %d)>" % self._fd
299
300    @classmethod
301    def fromHandle(cls, handle):
302        """
303        Create a node from a file handle.
304
305        @param handle: A file-like object backed by a file descriptor.
306        """
307
308        self = cls(loadDRM(), handle.fileno())
309        # We must keep the handle alive, and we cannot trust the caller to
310        # keep it alive for us.
311        self._handle = handle
312        return self
313
314    def version(self):
315        """
316        Obtain the version.
317        """
318
319        v = self._l.drmGetVersion(self._fd).contents
320        v._l = self._l
321        return v
322
323    def resources(self):
324        """
325        Obtain the modesetting resources.
326        """
327
328        resources_ptr = self._l.drmModeGetResources(self._fd)
329        if resources_ptr:
330            r = resources_ptr.contents
331            r._fd = self._fd
332            r._l = self._l
333            return r
334
335        return None
336
337
338def drmFromPath(path):
339    """
340    Given a DRM node path, open the corresponding node.
341
342    @param path: The path of the minor node to open.
343    """
344
345    handle = open(path)
346    return DRM.fromHandle(handle)
347
348
349def _bgrx24(i):
350    b = ord(next(i))
351    g = ord(next(i))
352    r = ord(next(i))
353    next(i)
354    return r, g, b
355
356
357def _screenshot(image, fb):
358    fb.map()
359    m = fb._map
360    lineLength = fb.width * fb.bpp // 8
361    pitch = fb.pitch
362    pixels = []
363
364    if fb.depth == 24:
365        unformat = _bgrx24
366    else:
367        raise RuntimeError("Couldn't unformat FB: %r" % fb)
368
369    for y in range(fb.height):
370        offset = y * pitch
371        m.seek(offset)
372        channels = m.read(lineLength)
373        ichannels = iter(channels)
374        for x in range(fb.width):
375            rgb = unformat(ichannels)
376            image.putpixel((x, y), rgb)
377
378    fb.unmap()
379
380    return pixels
381
382
383_drm = None
384
385def crtcScreenshot(crtc_id=None):
386    """
387    Take a screenshot, returning an image object.
388
389    @param crtc_id: The CRTC to screenshot.
390    """
391
392    global _drm
393    if not _drm:
394        paths = ["/dev/dri/" + n for n in os.listdir("/dev/dri")]
395        for p in paths:
396            d = drmFromPath(p)
397            if d.resources():
398                _drm = d
399                break
400
401    if _drm:
402        fb = _drm.resources().getCrtc(crtc_id).fb()
403        image = Image.new("RGB", (fb.width, fb.height))
404        pixels = _screenshot(image, fb)
405        return image
406
407    raise RuntimeError("Couldn't screenshot with DRM devices")
408