1"""A wrapper around the Generic Buffer Manager (GBM) library.
2
3Currently implements exactly the functions required to screenshot a frame
4buffer using DRM crtc info.
5"""
6from ctypes import *
7import drm
8from PIL import Image
9
10GBM_BO_IMPORT_FD = 0x5503
11GBM_BO_USE_SCANOUT = c_uint(1)
12GBM_BO_TRANSFER_READ = c_uint(1)
13GBM_MAX_PLANES = 4
14
15def __gbm_fourcc_code(a, b, c, d):
16    return ord(a) | (ord(b) << 8) | (ord(c) << 16) | (ord(d) << 24)
17
18GBM_FORMAT_ARGB8888 = __gbm_fourcc_code("A", "R", "2", "4")
19GBM_LIBRARIES = ["libgbm.so", "libgbm.so.1"]
20
21
22class gbm_import_fd_data(Structure):
23    _fields_ = [
24        ("fd", c_int),
25        ("width", c_uint),
26        ("height", c_uint),
27        ("stride", c_uint),
28        ("bo_format", c_uint),
29    ]
30
31
32class gbm_import_fd_planar_data(Structure):
33    _fields_ = [
34        ("fds", c_int * GBM_MAX_PLANES),
35        ("width", c_uint),
36        ("height", c_uint),
37        ("bo_format", c_uint),
38        ("strides", c_uint * GBM_MAX_PLANES),
39        ("offsets", c_uint * GBM_MAX_PLANES),
40        ("format_modifiers", c_ulonglong * GBM_MAX_PLANES),
41    ]
42
43
44class gbm_device(Structure):
45    """Opaque struct for GBM device.
46    """
47    pass
48
49
50class gbm_bo(Structure):
51    """Opaque struct for GBM buffer.
52    """
53    pass
54
55
56def loadGBM():
57    """Load and return a handle to libgbm.so.
58    """
59    l = None
60
61    for lib in GBM_LIBRARIES:
62        try:
63            l = cdll.LoadLibrary(lib)
64        except OSError:
65            l = None
66        if l is not None:
67            break
68
69    if l is None:
70        raise RuntimeError("Could not load GBM library.")
71        return None
72
73    l.gbm_create_device.argtypes = [c_int]
74    l.gbm_create_device.restype = POINTER(gbm_device)
75
76    l.gbm_device_destroy.argtypes = [POINTER(gbm_device)]
77    l.gbm_device_destroy.restype = None
78
79    l.gbm_bo_import.argtypes = [POINTER(gbm_device), c_uint, c_void_p, c_uint]
80    l.gbm_bo_import.restype = POINTER(gbm_bo)
81
82    l.gbm_bo_map.argtypes = [
83        POINTER(gbm_bo), c_uint, c_uint, c_uint, c_uint, c_uint,
84        POINTER(c_uint),
85        POINTER(c_void_p), c_size_t
86    ]
87    l.gbm_bo_map.restype = c_void_p
88
89    l.gbm_bo_unmap.argtypes = [POINTER(gbm_bo), c_void_p]
90    l.gbm_bo_unmap.restype = None
91
92    return l
93
94
95class GBMBuffer(object):
96    """A GBM buffer.
97    """
98
99    def __init__(self, library, buffer):
100        self._l = library
101        self._buffer = buffer
102
103    @classmethod
104    def fromFD(cls, device, fd, width, height, stride, bo_format, usage):
105        """Create/import a GBM Buffer Object from a file descriptor.
106
107        @param device: GBM device object.
108        @param fd: a file descriptor for the buffer to be imported.
109        @param width: buffer width in pixels.
110        @param height: buffer height in pixels.
111        @param stride: buffer pitch; number of pixels between sequential rows.
112        @param bo_format: pixel format fourcc code, e.g. GBM_FORMAT_ARGB8888.
113        @param usage: buffer usage flags, e.g. GBM_BO_USE_SCANOUT.
114        """
115        bo_data = gbm_import_fd_data()
116        bo_data.fd = fd
117        bo_data.width = width
118        bo_data.height = height
119        bo_data.stride = stride
120        bo_data.bo_format = bo_format
121        buffer = device._l.gbm_bo_import(device._device, GBM_BO_IMPORT_FD,
122                                         byref(bo_data), usage)
123        if buffer is None:
124            raise RuntimeError("gbm_bo_import() returned NULL")
125
126        self = cls(device._l, buffer)
127        return self
128
129    def map(self, x, y, width, height, flags, plane):
130        """Map buffer data into this user-space.
131        Returns (address, stride_bytes): void* start address for pixel array,
132        number of BYTES between sequental rows of pixels.
133        @param flags: The union of the GBM_BO_TRANSFER_* flags for this buffer.
134        """
135        self._map_p = c_void_p(0)
136        stride_out = c_uint(0)
137        if width == 0 or height == 0:
138            raise RuntimeError("Map width and/or height is 0")
139        map_p = self._l.gbm_bo_map(self._buffer, x, y, width, height, flags,
140                                   byref(stride_out), byref(self._map_p), plane)
141        if stride_out is 0:
142            raise RuntimeError("gbm_bo_map() stride is 0")
143        if map_p is 0:
144            raise RuntimeError("gbm_bo_map() returned NULL")
145        return map_p, stride_out
146
147    def unmap(self, map_data_p):
148        self._l.gbm_bo_unmap(self._buffer, map_data_p)
149
150
151class GBMDevice(object):
152    """A GBM device.
153    """
154
155    def __init__(self, library, handle):
156        self._l = library
157        self._handle = handle
158        self._device = library.gbm_create_device(self._handle)
159
160    def __del__(self):
161        if self._l:
162            self._l.gbm_device_destroy(self._device)
163
164    @classmethod
165    def fromHandle(cls, handle):
166        """Create a device object from an open file descriptor.
167        """
168        self = cls(loadGBM(), handle)
169        return self
170
171
172def _bgrx24(i):
173    b = i & 255
174    g = (i >> 8) & 255
175    r = (i >> 16) & 255
176    return r, g, b
177
178
179def _copyImage(image, map_ints, map_ints_pitch, unformat):
180    width, height = image.size
181    for y in range(height):
182        y_offset = y * map_ints_pitch
183        for x in range(width):
184            rgb = unformat(map_ints[y_offset + x])
185            image.putpixel((x, y), rgb)
186
187
188def crtcScreenshot(crtc_id=None):
189    """Take a screenshot, returning an image object.
190
191    @param crtc_id: The CRTC to screenshot.
192                    None for first found CRTC with mode set
193                    or "internal" for crtc connected to internal LCD
194                    or "external" for crtc connected to external display
195                    or "usb" "evdi" or "udl" for crtc with valid mode on evdi
196                    or udl display
197                    or DRM integer crtc_id
198    """
199    crtc = drm.getCrtc(crtc_id)
200    if crtc is not None:
201        device = GBMDevice.fromHandle(drm._drm._fd)
202        framebuffer = crtc.fb()
203        # TODO(djmk): The buffer format is hardcoded to ARGB8888, we should fix
204        # this to query for the frambuffer's format instead.
205        format_hardcode = GBM_FORMAT_ARGB8888
206
207        bo = GBMBuffer.fromFD(device,
208                              framebuffer.getFD(), framebuffer.width,
209                              framebuffer.height, framebuffer.pitch,
210                              format_hardcode, GBM_BO_USE_SCANOUT)
211        map_void_p, stride_bytes = bo.map(0, 0, framebuffer.width,
212                                          framebuffer.height,
213                                          GBM_BO_TRANSFER_READ, 0)
214        stride_pixels = stride_bytes.value / 4
215        map_ints_type = c_int * (stride_pixels * framebuffer.height)
216        map_int_p = cast(map_void_p, POINTER(c_int))
217        addr = addressof(map_int_p.contents)
218        map_ints = map_ints_type.from_address(addr)
219        image = Image.new("RGB", (framebuffer.width, framebuffer.height))
220        _copyImage(image, map_ints, stride_pixels, _bgrx24)
221        bo.unmap(bo._map_p)
222        return image
223
224    raise RuntimeError(
225        "Unable to take screenshot. There may not be anything on the screen.")
226