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    def __del__(self):
104        if self._l:
105            self._l.gbm_bo_destroy(self._buffer)
106
107    @classmethod
108    def fromFD(cls, device, fd, width, height, stride, bo_format, usage):
109        """Create/import a GBM Buffer Object from a file descriptor.
110
111        @param device: GBM device object.
112        @param fd: a file descriptor for the buffer to be imported.
113        @param width: buffer width in pixels.
114        @param height: buffer height in pixels.
115        @param stride: buffer pitch; number of pixels between sequential rows.
116        @param bo_format: pixel format fourcc code, e.g. GBM_FORMAT_ARGB8888.
117        @param usage: buffer usage flags, e.g. GBM_BO_USE_SCANOUT.
118        """
119        bo_data = gbm_import_fd_data()
120        bo_data.fd = fd
121        bo_data.width = width
122        bo_data.height = height
123        bo_data.stride = stride
124        bo_data.bo_format = bo_format
125        buffer = device._l.gbm_bo_import(device._device, GBM_BO_IMPORT_FD,
126                                         byref(bo_data), usage)
127        if buffer is None:
128            raise RuntimeError("gbm_bo_import() returned NULL")
129
130        self = cls(device._l, buffer)
131        return self
132
133    def map(self, x, y, width, height, flags, plane):
134        """Map buffer data into this user-space.
135        Returns (address, stride_bytes): void* start address for pixel array,
136        number of BYTES between sequental rows of pixels.
137        @param flags: The union of the GBM_BO_TRANSFER_* flags for this buffer.
138        """
139        self._map_p = c_void_p(0)
140        stride_out = c_uint(0)
141        if width == 0 or height == 0:
142            raise RuntimeError("Map width and/or height is 0")
143        map_p = self._l.gbm_bo_map(self._buffer, x, y, width, height, flags,
144                                   byref(stride_out), byref(self._map_p), plane)
145        if stride_out is 0:
146            raise RuntimeError("gbm_bo_map() stride is 0")
147        if map_p is 0:
148            raise RuntimeError("gbm_bo_map() returned NULL")
149        return map_p, stride_out
150
151    def unmap(self, map_data_p):
152        self._l.gbm_bo_unmap(self._buffer, map_data_p)
153
154
155class GBMDevice(object):
156    """A GBM device.
157    """
158
159    def __init__(self, library, handle):
160        self._l = library
161        self._handle = handle
162        self._device = library.gbm_create_device(self._handle)
163
164    def __del__(self):
165        if self._l:
166            self._l.gbm_device_destroy(self._device)
167
168    @classmethod
169    def fromHandle(cls, handle):
170        """Create a device object from an open file descriptor.
171        """
172        self = cls(loadGBM(), handle)
173        return self
174
175
176def _bgrx24(i):
177    b = i & 255
178    g = (i >> 8) & 255
179    r = (i >> 16) & 255
180    return r, g, b
181
182
183def crtcScreenshot(crtc_id=None):
184    """Take a screenshot, returning an image object.
185
186    @param crtc_id: The CRTC to screenshot.
187                    None for first found CRTC with mode set
188                    or "internal" for crtc connected to internal LCD
189                    or "external" for crtc connected to external display
190                    or "usb" "evdi" or "udl" for crtc with valid mode on evdi
191                    or udl display
192                    or DRM integer crtc_id
193    """
194    crtc = drm.getCrtc(crtc_id)
195    if crtc is not None:
196        device = GBMDevice.fromHandle(drm._drm._fd)
197        framebuffer = crtc.fb()
198        # TODO(djmk): The buffer format is hardcoded to ARGB8888, we should fix
199        # this to query for the frambuffer's format instead.
200        format_hardcode = GBM_FORMAT_ARGB8888
201
202        bo = GBMBuffer.fromFD(device,
203                              framebuffer.getFD(), framebuffer.width,
204                              framebuffer.height, framebuffer.pitch,
205                              format_hardcode, GBM_BO_USE_SCANOUT)
206        map_void_p, stride_bytes = bo.map(0, 0, framebuffer.width,
207                                          framebuffer.height,
208                                          GBM_BO_TRANSFER_READ, 0)
209        map_bytes = stride_bytes.value * framebuffer.height
210
211        # Create a Python Buffer object which references (but does not own) the
212        # memory.
213        buffer_from_memory = pythonapi.PyBuffer_FromMemory
214        buffer_from_memory.restype = py_object
215        buffer_from_memory.argtypes = [c_void_p, c_ssize_t]
216        map_buffer = buffer_from_memory(map_void_p, map_bytes)
217
218        # Make a copy of the bytes. Doing this is faster than the conversion,
219        # and is more likely to capture a consistent snapshot of the framebuffer
220        # contents, as a process may be writing to it.
221        buffer_bytes = bytes(map_buffer)
222
223        # Load the image, converting from the BGRX format to a PIL Image in RGB
224        # form. As the conversion is implemented by PIL as C code, this
225        # conversion is much faster than calling _bgrx24().
226        image = Image.fromstring(
227                'RGB', (framebuffer.width, framebuffer.height), buffer_bytes,
228                'raw', 'BGRX', stride_bytes.value, 1)
229        bo.unmap(bo._map_p)
230	del bo
231        return image
232
233    raise RuntimeError(
234        "Unable to take screenshot. There may not be anything on the screen.")
235
236
237def crtcGetPixel(x, y, crtc_id=None):
238    """Get a pixel from the specified screen, returning a rgb tuple.
239
240    @param x: The x-coordinate of the desired pixel.
241    @param y: The y-coordinate of the desired pixel.
242    @param crtc_id: The CRTC to get the pixel from.
243                    None for first found CRTC with mode set
244                    or "internal" for crtc connected to internal LCD
245                    or "external" for crtc connected to external display
246                    or "usb" "evdi" or "udl" for crtc with valid mode on evdi
247                    or udl display
248                    or DRM integer crtc_id
249    """
250    crtc = drm.getCrtc(crtc_id)
251    if crtc is None:
252        raise RuntimeError(
253            "Unable to get pixel. There may not be anything on the screen.")
254    device = GBMDevice.fromHandle(drm._drm._fd)
255    framebuffer = crtc.fb()
256    bo = GBMBuffer.fromFD(device,
257                          framebuffer.getFD(), framebuffer.width,
258                          framebuffer.height, framebuffer.pitch,
259                          GBM_FORMAT_ARGB8888, GBM_BO_USE_SCANOUT)
260    map_void_p, _ = bo.map(int(x), int(y), 1, 1, GBM_BO_TRANSFER_READ, 0)
261    map_int_p = cast(map_void_p, POINTER(c_int))
262    pixel = _bgrx24(map_int_p[0])
263    bo.unmap(bo._map_p)
264    return pixel
265