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 exceptions
16import mmap
17import os
18import subprocess
19
20from PIL import Image
21
22# drmModeConnection enum
23DRM_MODE_CONNECTED         = 1
24DRM_MODE_DISCONNECTED      = 2
25DRM_MODE_UNKNOWNCONNECTION = 3
26
27DRM_MODE_CONNECTOR_Unknown     = 0
28DRM_MODE_CONNECTOR_VGA         = 1
29DRM_MODE_CONNECTOR_DVII        = 2
30DRM_MODE_CONNECTOR_DVID        = 3
31DRM_MODE_CONNECTOR_DVIA        = 4
32DRM_MODE_CONNECTOR_Composite   = 5
33DRM_MODE_CONNECTOR_SVIDEO      = 6
34DRM_MODE_CONNECTOR_LVDS        = 7
35DRM_MODE_CONNECTOR_Component   = 8
36DRM_MODE_CONNECTOR_9PinDIN     = 9
37DRM_MODE_CONNECTOR_DisplayPort = 10
38DRM_MODE_CONNECTOR_HDMIA       = 11
39DRM_MODE_CONNECTOR_HDMIB       = 12
40DRM_MODE_CONNECTOR_TV          = 13
41DRM_MODE_CONNECTOR_eDP         = 14
42DRM_MODE_CONNECTOR_VIRTUAL     = 15
43DRM_MODE_CONNECTOR_DSI         = 16
44
45
46class DrmVersion(Structure):
47    """
48    The version of a DRM node.
49    """
50
51    _fields_ = [
52        ("version_major", c_int),
53        ("version_minor", c_int),
54        ("version_patchlevel", c_int),
55        ("name_len", c_int),
56        ("name", c_char_p),
57        ("date_len", c_int),
58        ("date", c_char_p),
59        ("desc_len", c_int),
60        ("desc", c_char_p),
61    ]
62
63    _l = None
64
65    def __repr__(self):
66        return "%s %d.%d.%d (%s) (%s)" % (self.name,
67                                          self.version_major,
68                                          self.version_minor,
69                                          self.version_patchlevel,
70                                          self.desc,
71                                          self.date,)
72
73    def __del__(self):
74        if self._l:
75            self._l.drmFreeVersion(self)
76
77
78class DrmModeResources(Structure):
79    """
80    Resources associated with setting modes on a DRM node.
81    """
82
83    _fields_ = [
84        ("count_fbs", c_int),
85        ("fbs", POINTER(c_uint)),
86        ("count_crtcs", c_int),
87        ("crtcs", POINTER(c_uint)),
88        ("count_connectors", c_int),
89        ("connectors", POINTER(c_uint)),
90        ("count_encoders", c_int),
91        ("encoders", POINTER(c_uint)),
92        ("min_width", c_int),
93        ("max_width", c_int),
94        ("min_height", c_int),
95        ("max_height", c_int),
96    ]
97
98    _fd = None
99    _l = None
100
101    def __repr__(self):
102        return "<DRM mode resources>"
103
104    def __del__(self):
105        if self._l:
106            self._l.drmModeFreeResources(self)
107
108    def _wakeup_screen(self):
109        """
110        Send a synchronous dbus message to power on screen.
111        """
112        # Get and process reply to make this synchronous.
113        subprocess.check_output([
114            "dbus-send", "--type=method_call", "--system", "--print-reply",
115            "--dest=org.chromium.PowerManager", "/org/chromium/PowerManager",
116            "org.chromium.PowerManager.HandleUserActivity", "int32:0"
117        ])
118
119    def getValidCrtc(self):
120        for i in xrange(0, self.count_crtcs):
121            crtc_id = self.crtcs[i]
122            crtc = self._l.drmModeGetCrtc(self._fd, crtc_id).contents
123            if crtc.mode_valid:
124                return crtc
125        return None
126
127    def getCrtc(self, crtc_id):
128        """
129        Obtain the CRTC at a given index.
130
131        @param crtc_id: The CRTC to get.
132        """
133        if crtc_id:
134            return self._l.drmModeGetCrtc(self._fd, crtc_id).contents
135        return self.getValidCrtc()
136
137    def getCrtcRobust(self, crtc_id=None):
138        crtc = self.getCrtc(crtc_id)
139        if crtc is None:
140            self._wakeup_screen()
141            crtc = self.getCrtc(crtc_id)
142        if crtc is not None:
143            crtc._fd = self._fd
144            crtc._l = self._l
145        return crtc
146
147
148class DrmModeModeInfo(Structure):
149    """
150    A DRM modesetting mode info.
151    """
152
153    _fields_ = [
154        ("clock", c_uint),
155        ("hdisplay", c_ushort),
156        ("hsync_start", c_ushort),
157        ("hsync_end", c_ushort),
158        ("htotal", c_ushort),
159        ("hskew", c_ushort),
160        ("vdisplay", c_ushort),
161        ("vsync_start", c_ushort),
162        ("vsync_end", c_ushort),
163        ("vtotal", c_ushort),
164        ("vscan", c_ushort),
165        ("vrefresh", c_uint),
166        ("flags", c_uint),
167        ("type", c_uint),
168        ("name", c_char * 32),
169    ]
170
171
172class DrmModeCrtc(Structure):
173    """
174    A DRM modesetting CRTC.
175    """
176
177    _fields_ = [
178        ("crtc_id", c_uint),
179        ("buffer_id", c_uint),
180        ("x", c_uint),
181        ("y", c_uint),
182        ("width", c_uint),
183        ("height", c_uint),
184        ("mode_valid", c_int),
185        ("mode", DrmModeModeInfo),
186        ("gamma_size", c_int),
187    ]
188
189    _fd = None
190    _l = None
191
192    def __repr__(self):
193        return "<CRTC (%d)>" % self.crtc_id
194
195    def __del__(self):
196        if self._l:
197            self._l.drmModeFreeCrtc(self)
198
199    def hasFb(self):
200        """
201        Whether this CRTC has an associated framebuffer.
202        """
203
204        return self.buffer_id != 0
205
206    def fb(self):
207        """
208        Obtain the framebuffer, if one is associated.
209        """
210
211        if self.hasFb():
212            fb = self._l.drmModeGetFB(self._fd, self.buffer_id).contents
213            fb._fd = self._fd
214            fb._l = self._l
215            return fb
216        else:
217            raise RuntimeError("CRTC %d doesn't have a framebuffer!" %
218                               self.crtc_id)
219
220
221class DrmModeEncoder(Structure):
222    """
223    A DRM modesetting encoder.
224    """
225
226    _fields_ = [
227        ("encoder_id", c_uint),
228        ("encoder_type", c_uint),
229        ("crtc_id", c_uint),
230        ("possible_crtcs", c_uint),
231        ("possible_clones", c_uint),
232    ]
233
234    _fd = None
235    _l = None
236
237    def __repr__(self):
238        return "<Encoder (%d)>" % self.encoder_id
239
240    def __del__(self):
241        if self._l:
242            self._l.drmModeFreeEncoder(self)
243
244
245class DrmModeConnector(Structure):
246    """
247    A DRM modesetting connector.
248    """
249
250    _fields_ = [
251        ("connector_id", c_uint),
252        ("encoder_id", c_uint),
253        ("connector_type", c_uint),
254        ("connector_type_id", c_uint),
255        ("connection", c_uint), # drmModeConnection enum
256        ("mmWidth", c_uint),
257        ("mmHeight", c_uint),
258        ("subpixel", c_uint), # drmModeSubPixel enum
259        ("count_modes", c_int),
260        ("modes", POINTER(DrmModeModeInfo)),
261        ("count_propts", c_int),
262        ("props", POINTER(c_uint)),
263        ("prop_values", POINTER(c_ulonglong)),
264        ("count_encoders", c_int),
265        ("encoders", POINTER(c_uint)),
266    ]
267
268    _fd = None
269    _l = None
270
271    def __repr__(self):
272        return "<Connector (%d)>" % self.connector_id
273
274    def __del__(self):
275        if self._l:
276            self._l.drmModeFreeConnector(self)
277
278    def isInternal(self):
279        return (self.connector_type == DRM_MODE_CONNECTOR_LVDS or
280                self.connector_type == DRM_MODE_CONNECTOR_eDP or
281                self.connector_type == DRM_MODE_CONNECTOR_DSI)
282
283    def isConnected(self):
284        return self.connection == DRM_MODE_CONNECTED
285
286
287class drm_mode_map_dumb(Structure):
288    """
289    Request a mapping of a modesetting buffer.
290
291    The map will be "dumb;" it will be accessible via mmap() but very slow.
292    """
293
294    _fields_ = [
295        ("handle", c_uint),
296        ("pad", c_uint),
297        ("offset", c_ulonglong),
298    ]
299
300
301# This constant is not defined in any one header; it is the pieced-together
302# incantation for the ioctl that performs dumb mappings. I would love for this
303# to not have to be here, but it can't be imported from any header easily.
304DRM_IOCTL_MODE_MAP_DUMB = 0xc01064b3
305
306
307class DrmModeFB(Structure):
308    """
309    A DRM modesetting framebuffer.
310    """
311
312    _fields_ = [
313        ("fb_id", c_uint),
314        ("width", c_uint),
315        ("height", c_uint),
316        ("pitch", c_uint),
317        ("bpp", c_uint),
318        ("depth", c_uint),
319        ("handle", c_uint),
320    ]
321
322    _l = None
323    _map = None
324
325    def __repr__(self):
326        s = "<Framebuffer (%dx%d (pitch %d bytes), %d bits/pixel, depth %d)"
327        vitals = s % (self.width,
328                      self.height,
329                      self.pitch,
330                      self.bpp,
331                      self.depth,)
332        if self._map:
333            tail = " (mapped)>"
334        else:
335            tail = ">"
336        return vitals + tail
337
338    def __del__(self):
339        if self._l:
340            self._l.drmModeFreeFB(self)
341
342    def map(self, size):
343        """
344        Map the framebuffer.
345        """
346
347        if self._map:
348            return
349
350        mapDumb = drm_mode_map_dumb()
351        mapDumb.handle = self.handle
352
353        rv = self._l.drmIoctl(self._fd, DRM_IOCTL_MODE_MAP_DUMB,
354                              pointer(mapDumb))
355        if rv:
356            raise IOError(rv, os.strerror(rv))
357
358        # mmap.mmap() has a totally different order of arguments in Python
359        # compared to C; check the documentation before altering this
360        # incantation.
361        self._map = mmap.mmap(self._fd,
362                              size,
363                              flags=mmap.MAP_SHARED,
364                              prot=mmap.PROT_READ,
365                              offset=mapDumb.offset)
366
367    def unmap(self):
368        """
369        Unmap the framebuffer.
370        """
371
372        if self._map:
373            self._map.close()
374            self._map = None
375
376
377def loadDRM():
378    """
379    Load a handle to libdrm.
380
381    In addition to loading, this function also configures the argument and
382    return types of functions.
383    """
384
385    l = None
386
387    try:
388        l = cdll.LoadLibrary("libdrm.so")
389    except OSError:
390        l = cdll.LoadLibrary("libdrm.so.2") # ubuntu doesn't have libdrm.so
391
392    l.drmGetVersion.argtypes = [c_int]
393    l.drmGetVersion.restype = POINTER(DrmVersion)
394
395    l.drmFreeVersion.argtypes = [POINTER(DrmVersion)]
396    l.drmFreeVersion.restype = None
397
398    l.drmModeGetResources.argtypes = [c_int]
399    l.drmModeGetResources.restype = POINTER(DrmModeResources)
400
401    l.drmModeFreeResources.argtypes = [POINTER(DrmModeResources)]
402    l.drmModeFreeResources.restype = None
403
404    l.drmModeGetCrtc.argtypes = [c_int, c_uint]
405    l.drmModeGetCrtc.restype = POINTER(DrmModeCrtc)
406
407    l.drmModeFreeCrtc.argtypes = [POINTER(DrmModeCrtc)]
408    l.drmModeFreeCrtc.restype = None
409
410    l.drmModeGetEncoder.argtypes = [c_int, c_uint]
411    l.drmModeGetEncoder.restype = POINTER(DrmModeEncoder)
412
413    l.drmModeFreeEncoder.argtypes = [POINTER(DrmModeEncoder)]
414    l.drmModeFreeEncoder.restype = None
415
416    l.drmModeGetConnector.argtypes = [c_int, c_uint]
417    l.drmModeGetConnector.restype = POINTER(DrmModeConnector)
418
419    l.drmModeFreeConnector.argtypes = [POINTER(DrmModeConnector)]
420    l.drmModeFreeConnector.restype = None
421
422    l.drmModeGetFB.argtypes = [c_int, c_uint]
423    l.drmModeGetFB.restype = POINTER(DrmModeFB)
424
425    l.drmModeFreeFB.argtypes = [POINTER(DrmModeFB)]
426    l.drmModeFreeFB.restype = None
427
428    l.drmIoctl.argtypes = [c_int, c_ulong, c_voidp]
429    l.drmIoctl.restype = c_int
430
431    return l
432
433
434class DRM(object):
435    """
436    A DRM node.
437    """
438
439    def __init__(self, library, fd):
440        self._l = library
441        self._fd = fd
442
443    def __repr__(self):
444        return "<DRM (FD %d)>" % self._fd
445
446    @classmethod
447    def fromHandle(cls, handle):
448        """
449        Create a node from a file handle.
450
451        @param handle: A file-like object backed by a file descriptor.
452        """
453
454        self = cls(loadDRM(), handle.fileno())
455        # We must keep the handle alive, and we cannot trust the caller to
456        # keep it alive for us.
457        self._handle = handle
458        return self
459
460    def version(self):
461        """
462        Obtain the version.
463        """
464
465        v = self._l.drmGetVersion(self._fd).contents
466        v._l = self._l
467        return v
468
469    def resources(self):
470        """
471        Obtain the modesetting resources.
472        """
473
474        resources_ptr = self._l.drmModeGetResources(self._fd)
475        if resources_ptr:
476            r = resources_ptr.contents
477            r._fd = self._fd
478            r._l = self._l
479            return r
480
481        return None
482
483    def getCrtc(self, crtc_id):
484        c_ptr = self._l.drmModeGetCrtc(self._fd, crtc_id)
485        if c_ptr:
486            c = c_ptr.contents
487            c._fd = self._fd
488            c._l = self._l
489            return c
490
491        return None
492
493    def getEncoder(self, encoder_id):
494        e_ptr = self._l.drmModeGetEncoder(self._fd, encoder_id)
495        if e_ptr:
496            e = e_ptr.contents
497            e._fd = self._fd
498            e._l = self._l
499            return e
500
501        return None
502
503    def getConnector(self, connector_id):
504        c_ptr = self._l.drmModeGetConnector(self._fd, connector_id)
505        if c_ptr:
506            c = c_ptr.contents
507            c._fd = self._fd
508            c._l = self._l
509            return c
510
511        return None
512
513
514
515def drmFromPath(path):
516    """
517    Given a DRM node path, open the corresponding node.
518
519    @param path: The path of the minor node to open.
520    """
521
522    handle = open(path)
523    return DRM.fromHandle(handle)
524
525
526def _bgrx24(i):
527    b = ord(next(i))
528    g = ord(next(i))
529    r = ord(next(i))
530    next(i)
531    return r, g, b
532
533
534def _copyImageBlocklinear(image, fb, unformat):
535    gobPitch = 64
536    gobHeight = 128
537    while gobHeight > 8 and gobHeight >= 2 * fb.height:
538        gobHeight //= 2
539    gobSize = gobPitch * gobHeight
540    gobWidth = gobPitch // (fb.bpp // 8)
541
542    gobCountX = (fb.pitch + gobPitch - 1) // gobPitch
543    gobCountY = (fb.height + gobHeight - 1) // gobHeight
544    fb.map(gobCountX * gobCountY * gobSize)
545    m = fb._map
546
547    offset = 0
548    for gobY in range(gobCountY):
549        gobTop = gobY * gobHeight
550        for gobX in range(gobCountX):
551            m.seek(offset)
552            gob = m.read(gobSize)
553            iterGob = iter(gob)
554            gobLeft = gobX * gobWidth
555            for i in range(gobWidth * gobHeight):
556                rgb = unformat(iterGob)
557                x = gobLeft + (((i >> 3) & 8) | ((i >> 1) & 4) | (i & 3))
558                y = gobTop + ((i >> 7 << 3) | ((i >> 3) & 6) | ((i >> 2) & 1))
559                if x < fb.width and y < fb.height:
560                    image.putpixel((x, y), rgb)
561            offset += gobSize
562    fb.unmap()
563
564
565def _copyImageLinear(image, fb, unformat):
566    fb.map(fb.pitch * fb.height)
567    m = fb._map
568    pitch = fb.pitch
569    lineLength = fb.width * fb.bpp // 8
570    for y in range(fb.height):
571        offset = y * pitch
572        m.seek(offset)
573        channels = m.read(lineLength)
574        ichannels = iter(channels)
575        for x in range(fb.width):
576            rgb = unformat(ichannels)
577            image.putpixel((x, y), rgb)
578    fb.unmap()
579
580
581def _screenshot(drm, image, fb):
582    if fb.depth == 24:
583        unformat = _bgrx24
584    else:
585        raise RuntimeError("Couldn't unformat FB: %r" % fb)
586
587    if drm.version().name == "tegra":
588        _copyImageBlocklinear(image, fb, unformat)
589    else:
590        _copyImageLinear(image, fb, unformat)
591
592
593_drm = None
594
595
596def crtcScreenshot(crtc_id=None):
597    """
598    Take a screenshot, returning an image object.
599
600    @param crtc_id: The CRTC to screenshot.
601                    None for first found CRTC with mode set
602                    or "internal" for crtc connected to internal LCD
603                    or "external" for crtc connected to external display
604                    or "usb" "evdi" or "udl" for crtc with valid mode on evdi or
605                    udl display
606                    or DRM integer crtc_id
607    """
608
609    global _drm
610
611    if not _drm:
612        paths = [
613            "/dev/dri/" + n
614            for n in filter(lambda x: x.startswith("card"),
615                            os.listdir("/dev/dri"))
616        ]
617
618        if crtc_id == "usb" or crtc_id == "evdi" or crtc_id == "udl":
619            for p in paths:
620                d = drmFromPath(p)
621                v = d.version()
622
623                if crtc_id == v.name:
624                    _drm = d
625                    break
626
627                if crtc_id == "usb" and (v.name == "evdi" or v.name == "udl"):
628                    _drm = d
629                    break
630
631        elif crtc_id == "internal" or crtc_id == "external":
632            internal = crtc_id == "internal"
633            for p in paths:
634                d = drmFromPath(p)
635                if d.resources() is None:
636                    continue
637                if d.resources() and d.resources().count_connectors > 0:
638                    for c in xrange(0, d.resources().count_connectors):
639                        connector = d.getConnector(d.resources().connectors[c])
640                        if (internal == connector.isInternal()
641                            and connector.isConnected()
642                            and connector.encoder_id != 0):
643                            e = d.getEncoder(connector.encoder_id)
644                            crtc = d.getCrtc(e.crtc_id)
645                            if crtc.mode_valid:
646                                crtc_id = crtc.crtc_id
647                                _drm = d
648                                break
649                if _drm:
650                    break
651
652        elif crtc_id is None or crtc_id == 0:
653            for p in paths:
654                d = drmFromPath(p)
655                if d.resources() is None:
656                    continue
657                for c in xrange(0, d.resources().count_crtcs):
658                    crtc = d.getCrtc(d.resources().crtcs[c])
659                    if crtc.mode_valid:
660                        crtc_id = d.resources().crtcs[c]
661                        _drm = d
662                        break
663                if _drm:
664                    break
665
666        else:
667            for p in paths:
668                d = drmFromPath(p)
669                if d.resources() is None:
670                    continue
671                for c in xrange(0, d.resources().count_crtcs):
672                    if crtc_id == d.resources().crtcs[c]:
673                        _drm = d
674                        break
675                if _drm:
676                    break
677
678    if _drm:
679        crtc = _drm.resources().getCrtcRobust(crtc_id)
680        if crtc is not None:
681            framebuffer = crtc.fb()
682            image = Image.new("RGB", (framebuffer.width, framebuffer.height))
683            _screenshot(_drm, image, framebuffer)
684            return image
685
686    raise RuntimeError(
687        "Unable to take screenshot. There may not be anything on the screen.")
688