1# Copyright 2016 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""This module provides the utilities for chameleon streaming server usage.
6
7Sampe Code for dump realtime video frame:
8  stream = ChameleonStreamServer(IP)
9  stream.reset_video_session()
10
11  chameleon_proxy.StartCapturingVideo(port)
12  stream.dump_realtime_video_frame(False, RealtimeMode.BestEffort)
13  while True:
14    video_frame = stream.receive_realtime_video_frame()
15    if not video_frame:
16        break
17    (frame_number, width, height, channel, data) = video_frame
18    image = Image.fromstring('RGB', (width, height), data)
19    image.save('%d.bmp' % frame_number)
20
21Sampe Code for dump realtime audio page:
22  stream = ChameleonStreamServer(IP)
23  stream.reset_audio_session()
24
25  chameleon_proxy.StartCapturingAudio(port)
26  stream.dump_realtime_audio_page(RealtimeMode.BestEffort)
27  f = open('audio.raw',  'w')
28  while True:
29    audio_page = stream.receive_realtime_audio_page()
30    if not audio_page:
31        break
32    (page_count, data) = audio_page
33    f.write(data)
34
35"""
36
37import logging
38import socket
39from struct import calcsize, pack, unpack
40
41
42CHAMELEON_STREAN_SERVER_PORT = 9994
43SUPPORT_MAJOR_VERSION = 1
44SUPPORT_MINOR_VERSION = 0
45
46
47class StreamServerVersionError(Exception):
48    """Version is not compatible between client and server."""
49    pass
50
51
52class ErrorCode(object):
53    """Error codes of response from the stream server."""
54    OK = 0
55    NonSupportCommand = 1
56    Argument = 2
57    RealtimeStreamExists = 3
58    VideoMemoryOverflowStop = 4
59    VideoMemoryOverflowDrop = 5
60    AudioMemoryOverflowStop = 6
61    AudioMemoryOverflowDrop = 7
62    MemoryAllocFail = 8
63
64
65class RealtimeMode(object):
66    """Realtime mode of dumping data."""
67    # Stop dump when memory overflow
68    StopWhenOverflow = 1
69
70    # Drop data when memory overflow
71    BestEffort = 2
72
73    # Strings used for logging.
74    LogStrings = ['None', 'Stop when overflow', 'Best effort']
75
76
77class ChameleonStreamServer(object):
78    """
79    This class provides easy-to-use APIs to access the stream server.
80
81    """
82
83    # Main message types.
84    _REQUEST_TYPE = 0
85    _RESPONSE_TYPE = 1
86    _DATA_TYPE = 2
87
88    # Message types.
89    _Reset = 0
90    _GetVersion = 1
91    _ConfigVideoStream = 2
92    _ConfigShrinkVideoStream = 3
93    _DumpVideoFrame = 4
94    _DumpRealtimeVideoFrame = 5
95    _StopDumpVideoFrame = 6
96    _DumpRealtimeAudioPage = 7
97    _StopDumpAudioPage = 8
98
99    _PACKET_HEAD_SIZE = 8
100
101    # uint16 type, uint16 error_code, uint32 length.
102    packet_head_struct = '!HHL'
103    # uint8 major, uint8 minor.
104    version_struct = '!BB'
105    # unt16 screen_width, uint16 screen_height.
106    config_video_stream_struct = '!HH'
107    # uint8 shrink_width, uint8 shrink_height.
108    config_shrink_video_stream_struct = '!BB'
109    # uint32 memory_address1, uint32 memory_address2, uint16 number_of_frames.
110    dump_video_frame_struct = '!LLH'
111    # uint8 is_dual, uint8 mode.
112    dump_realtime_video_frame_struct = '!BB'
113    # uint32 frame_number, uint16 width, uint16 height, uint8 channel,
114    # uint8 padding[3]
115    video_frame_data_struct = '!LHHBBBB'
116    # uint8 mode.
117    dump_realtime_audio_page_struct = '!B'
118    # uint32 page_count.
119    audio_page_data_struct = '!L'
120
121    def __init__(self, hostname, port=CHAMELEON_STREAN_SERVER_PORT):
122        """Constructs a ChameleonStreamServer.
123
124        @param hostname: Hostname of stream server.
125        @param port: Port number the stream server is listening on.
126
127        """
128        self._video_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
129        self._audio_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
130        self._hostname = hostname
131        self._port = port
132        # Used for non-realtime dump video frames.
133        self._remain_frame_count = 0
134        self._is_realtime_video = False
135        self._is_realtime_audio = False
136
137    def _get_request_type(self, message):
138        """Get the request type of the message."""
139        return (self._REQUEST_TYPE << 8) | message
140
141    def _get_response_type(self, message):
142        """Get the response type of the message."""
143        return (self._RESPONSE_TYPE << 8) | message
144
145    def _is_data_type(self, message):
146        """Check if the message type is data type."""
147        return (self._DATA_TYPE << 8) & message
148
149    def _receive_whole_packet(self, sock):
150        """Receive one whole packet, contains packet head and content.
151
152        @param sock: Which socket to be used.
153
154        @return: A tuple with 4 elements contains message_type, error code,
155        length and content.
156
157        """
158        # receive packet header
159        data = sock.recv(self._PACKET_HEAD_SIZE)
160        if not data:
161            return None
162
163        while len(data) != self._PACKET_HEAD_SIZE:
164            remain_length = self._PACKET_HEAD_SIZE - len(data)
165            recv_content = sock.recv(remain_length)
166            data += recv_content
167
168        message_type, error_code, length = unpack(self.packet_head_struct, data)
169
170        # receive content
171        content = ''
172        remain_length = length
173        while remain_length:
174            recv_content = sock.recv(remain_length)
175            if not recv_content:
176                return None
177            remain_length -= len(recv_content)
178            content += recv_content
179
180        if error_code != ErrorCode.OK:
181            logging.warn('Receive error code %d, %r', error_code, content)
182
183        return (message_type, error_code, length, content)
184
185    def _send_and_receive(self, packet, sock, check_error=True):
186        """Send packet to server and receive response from server.
187
188        @param packet: The packet to be sent.
189        @param sock: Which socket to be used.
190        @param check_error: Check the error code. If this is True, this function
191        will check the error code from response and raise exception if the error
192        code is not OK.
193
194        @return: The response packet from server. A tuple with 4 elements
195        contains message_type, error code, length and content.
196
197        @raise ValueError if check_error and error code is not OK.
198
199        """
200        sock.send(packet)
201        packet = self._receive_whole_packet(sock)
202        if check_error:
203            (_, error_code, _, _) = packet
204            if error_code != ErrorCode.OK:
205                raise ValueError('Error code is not OK')
206
207        return packet
208
209    def _generate_packet_head(self, message_type, length):
210        """Generate packet head with request message.
211
212        @param message_type: Message type.
213        @param length: The length in the head.
214
215        @return: The packet head content.
216
217        """
218        head = pack(self.packet_head_struct,
219                    self._get_request_type(message_type),
220                    ErrorCode.OK, length)
221        return head
222
223    def _generate_config_video_stream_packet(self, width, height):
224        """Generate config video stream whole packet.
225
226        @param width: The screen width of the video frame by pixel per channel.
227        @param height: The screen height of the video frame by pixel per
228                       channel.
229
230        @return: The whole packet content.
231
232        """
233        content = pack(self.config_video_stream_struct, width, height)
234        head = self._generate_packet_head(self._ConfigVideoStream, len(content))
235        return head + content
236
237    def _generate_config_shrink_video_stream_packet(self, shrink_width,
238                                                    shrink_height):
239        """Generate config shrink video stream whole packet.
240
241        @param shrink_width: Shrink (shrink_width+1) pixels to 1 pixel when do
242                             video dump.
243        @param shrink_height: Shrink (shrink_height+1) to 1 height when do video
244                              dump.
245
246        @return: The whole packet content.
247
248        """
249        content = pack(self.config_shrink_video_stream_struct, shrink_width,
250                       shrink_height)
251        head = self._generate_packet_head(self._ConfigShrinkVideoStream,
252                                          len(content))
253        return head + content
254
255    def _generate_dump_video_stream_packet(self, count, address1, address2):
256        """Generate dump video stream whole packet.
257
258        @param count: Specify number of video frames for dumping.
259        @param address1: Dump memory address1.
260        @param address2: Dump memory address2. If it is 0. It means we only dump
261                         from address1.
262
263        @return: The whole packet content.
264
265        """
266        content = pack(self.dump_video_frame_struct, address1, address2, count)
267        head = self._generate_packet_head(self._DumpVideoFrame, len(content))
268        return head + content
269
270    def _generate_dump_realtime_video_stream_packet(self, is_dual, mode):
271        """Generate dump realtime video stream whole packet.
272
273        @param is_dual: False: means only dump from channel1,
274                        True: means dump from dual channels.
275        @param mode: The values of RealtimeMode.
276
277        @return: The whole packet content.
278
279        """
280        content = pack(self.dump_realtime_video_frame_struct, is_dual, mode)
281        head = self._generate_packet_head(self._DumpRealtimeVideoFrame,
282                                          len(content))
283        return head + content
284
285    def _generate_dump_realtime_audio_stream_packet(self, mode):
286        """Generate dump realtime audio stream whole packet.
287
288        @param mode: The values of RealtimeMode.
289
290        @return: The whole packet content.
291
292        """
293        content = pack(self.dump_realtime_audio_page_struct, mode)
294        head = self._generate_packet_head(self._DumpRealtimeAudioPage,
295                                          len(content))
296        return head + content
297
298    def _receive_video_frame(self):
299        """Receive one video frame from server.
300
301        This function will assume it only can receive video frame data packet
302        from server. Unless the error code is not OK.
303
304        @return A tuple with error code on first element.
305                if error code is OK. A decoded values will be stored in a tuple.
306                (error_code, frame number, width, height, channel, data)
307                if error code is not OK. It will return a tuple with
308                (error code, content). The content is the error message from
309                server.
310
311        @raise ValueError if packet is not data packet.
312
313        """
314        (message, error_code, _, content) = self._receive_whole_packet(
315            self._video_sock)
316        if error_code != ErrorCode.OK:
317            return (error_code, content)
318
319        if not self._is_data_type(message):
320            raise ValueError('Message is not data')
321
322        video_frame_head_size = calcsize(self.video_frame_data_struct)
323        frame_number, width, height, channel, _, _, _ = unpack(
324            self.video_frame_data_struct, content[:video_frame_head_size])
325        data = content[video_frame_head_size:]
326        return (error_code, frame_number, width, height, channel, data)
327
328    def _get_version(self):
329        """Get the version of the server.
330
331        @return A tuple with Major and Minor number of the server.
332
333        @raise ValueError if error code from response is not OK.
334
335        """
336        packet = self._generate_packet_head(self._GetVersion, 0)
337        (_, _, _, content) = self._send_and_receive(packet, self._video_sock)
338        return unpack(self.version_struct, content)
339
340    def _check_version(self):
341        """Check if this client is compatible with the server.
342
343        The major number must be the same and the minor number of the server
344        must larger then the client's.
345
346        @return Compatible or not
347
348        """
349        (major, minor) = self._get_version()
350        logging.debug('Major %d, minor %d', major, minor)
351        return major == SUPPORT_MAJOR_VERSION and minor >= SUPPORT_MINOR_VERSION
352
353    def connect(self):
354        """Connect to the server and check the compatibility."""
355        server_address = (self._hostname, self._port)
356        logging.info('connecting to %s:%s', self._hostname, self._port)
357        self._video_sock.connect(server_address)
358        self._audio_sock.connect(server_address)
359        if not self._check_version():
360            raise StreamServerVersionError()
361
362    def reset_video_session(self):
363        """Reset the video session.
364
365        @raise ValueError if error code from response is not OK.
366
367        """
368        logging.info('Reset session')
369        packet = self._generate_packet_head(self._Reset, 0)
370        self._send_and_receive(packet, self._video_sock)
371
372    def reset_audio_session(self):
373        """Reset the audio session.
374        For audio, we don't need to reset any thing.
375
376        """
377        pass
378
379    def config_video_stream(self, width, height):
380        """Configure the properties of the non-realtime video stream.
381
382        @param width: The screen width of the video frame by pixel per channel.
383        @param height: The screen height of the video frame by pixel per
384                       channel.
385
386        @raise ValueError if error code from response is not OK.
387
388        """
389        logging.info('Config video, width %d, height %d', width, height)
390        packet = self._generate_config_video_stream_packet(width, height)
391        self._send_and_receive(packet, self._video_sock)
392
393    def config_shrink_video_stream(self, shrink_width, shrink_height):
394        """Configure the shrink operation of the video frame dump.
395
396        @param shrink_width: Shrink (shrink_width+1) pixels to 1 pixel when do
397                             video dump. 0 means no shrink.
398        @param shrink_height: Shrink (shrink_height+1) to 1 height when do video
399                              dump. 0 means no shrink.
400
401        @raise ValueError if error code from response is not OK.
402
403        """
404        logging.info('Config shrink video, shirnk_width %d, shrink_height %d',
405                     shrink_width, shrink_height)
406        packet = self._generate_config_shrink_video_stream_packet(shrink_width,
407                                                                  shrink_height)
408        self._send_and_receive(packet, self._video_sock)
409
410    def dump_video_frame(self, count, address1, address2):
411        """Ask server to dump video frames.
412
413        User must use receive_video_frame() to receive video frames after
414        calling this API.
415
416        Sampe Code:
417            address = chameleon_proxy.GetCapturedFrameAddresses(0)
418            count = chameleon_proxy.GetCapturedFrameCount()
419            server.dump_video_frame(count, int(address), 0)
420            while True:
421                video_frame = server.receive_video_frame()
422                if not video_frame:
423                    break
424                (frame_number, width, height, channel, data) = video_frame
425                image = Image.fromstring('RGB', (width, height), data)
426                image.save('%s.bmp' % frame_number)
427
428        @param count: Specify number of video frames.
429        @param address1: Dump memory address1.
430        @param address2: Dump memory address2. If it is 0. It means we only dump
431                         from address1.
432
433        @raise ValueError if error code from response is not OK.
434
435        """
436        logging.info('dump video frame count %d, address1 0x%x, address2 0x%x',
437                     count, address1, address2)
438        packet = self._generate_dump_video_stream_packet(count, address1,
439                                                         address2)
440        self._send_and_receive(packet, self._video_sock)
441        self._remain_frame_count = count
442
443    def dump_realtime_video_frame(self, is_dual, mode):
444        """Ask server to dump realtime video frames.
445
446        User must use receive_realtime_video_frame() to receive video frames
447        after calling this API.
448
449        Sampe Code:
450            server.dump_realtime_video_frame(False,
451                                             RealtimeMode.StopWhenOverflow)
452            while True:
453                video_frame = server.receive_realtime_video_frame()
454                if not video_frame:
455                    break
456                (frame_number, width, height, channel, data) = video_frame
457                image = Image.fromstring('RGB', (width, height), data)
458                image.save('%s.bmp' % frame_number)
459
460        @param is_dual: False: means only dump from channel1,
461                        True: means dump from dual channels.
462        @param mode: The values of RealtimeMode.
463
464        @raise ValueError if error code from response is not OK.
465
466        """
467        logging.info('dump realtime video frame is_dual %d, mode %s', is_dual,
468                     RealtimeMode.LogStrings[mode])
469        packet = self._generate_dump_realtime_video_stream_packet(is_dual, mode)
470        self._send_and_receive(packet, self._video_sock)
471        self._is_realtime_video = True
472
473    def receive_video_frame(self):
474        """Receive one video frame from server after calling dump_video_frame().
475
476        This function will assume it only can receive video frame data packet
477        from server. Unless the error code is not OK.
478
479        @return A tuple with video frame information.
480                (frame number, width, height, channel, data)
481                None if error happens.
482
483        @raise ValueError if packet is not data packet.
484
485        """
486        if not self._remain_frame_count:
487            return None
488        self._remain_frame_count -= 1
489        frame_info = self._receive_video_frame()
490        if frame_info[0] != ErrorCode.OK:
491            self._remain_frame_count = 0
492            return None
493        return frame_info[1:]
494
495    def receive_realtime_video_frame(self):
496        """Receive one video frame from server after calling
497        dump_realtime_video_frame(). The video frame may be dropped if we use
498        BestEffort mode. We can detect it by the frame number.
499
500        This function will assume it only can receive video frame data packet
501        from server. Unless the error code is not OK.
502
503        @return A tuple with video frame information.
504                (frame number, width, height, channel, data)
505                None if error happens or no more frames.
506
507        @raise ValueError if packet is not data packet.
508
509        """
510        if not self._is_realtime_video:
511            return None
512
513        frame_info = self._receive_video_frame()
514        # We can still receive video frame for drop case.
515        while frame_info[0] == ErrorCode.VideoMemoryOverflowDrop:
516            frame_info = self._receive_video_frame()
517
518        if frame_info[0] != ErrorCode.OK:
519            return None
520
521        return frame_info[1:]
522
523    def stop_dump_realtime_video_frame(self):
524        """Ask server to stop dump realtime video frame."""
525        if not self._is_realtime_video:
526            return
527        packet = self._generate_packet_head(self._StopDumpVideoFrame, 0)
528        self._video_sock.send(packet)
529        # Drop video frames until receive _StopDumpVideoFrame response.
530        while True:
531            (message, _, _, _) = self._receive_whole_packet(self._video_sock)
532            if message == self._get_response_type(self._StopDumpVideoFrame):
533                break
534        self._is_realtime_video = False
535
536    def dump_realtime_audio_page(self, mode):
537        """Ask server to dump realtime audio pages.
538
539        User must use receive_realtime_audio_page() to receive audio pages
540        after calling this API.
541
542        Sampe Code for BestEffort:
543            server.dump_realtime_audio_page(RealtimeMode.kBestEffort)
544            f = open('audio.raw'), 'w')
545            while True:
546                audio_page = server.receive_realtime_audio_page()
547                if audio_page:
548                    break
549                (page_count, data) = audio_page
550                f.write(data)
551
552        @param mode: The values of RealtimeMode.
553
554        @raise ValueError if error code from response is not OK.
555
556        """
557        logging.info('dump realtime audio page mode %s',
558                     RealtimeMode.LogStrings[mode])
559        packet = self._generate_dump_realtime_audio_stream_packet(mode)
560        self._send_and_receive(packet, self._audio_sock)
561        self._is_realtime_audio = True
562
563    def receive_realtime_audio_page(self):
564        """Receive one audio page from server after calling
565        dump_realtime_audio_page(). The behavior is the same as
566        receive_realtime_video_frame(). The audio page may be dropped if we use
567        BestEffort mode. We can detect it by the page count.
568
569        This function will assume it only can receive audio page data packet
570        from server. Unless the error code is not OK.
571
572        @return A tuple with audio page information. (page count, data)
573                None if error happens or no more frames.
574
575        @raise ValueError if packet is not data packet.
576
577        """
578        if not self._is_realtime_audio:
579            return None
580        (message, error_code, _, content) = self._receive_whole_packet(
581            self._audio_sock)
582        # We can still receive audio page for drop case.
583        while error_code == ErrorCode.AudioMemoryOverflowDrop:
584            (message, error_code, _, content) = self._receive_whole_packet(
585                self._audio_sock)
586
587        if error_code != ErrorCode.OK:
588            return None
589        if not self._is_data_type(message):
590            raise ValueError('Message is not data')
591
592        page_count = unpack(self.audio_page_data_struct, content[:4])[0]
593        data = content[4:]
594        return (page_count, data)
595
596    def stop_dump_realtime_audio_page(self):
597        """Ask server to stop dump realtime audio page."""
598        if not self._is_realtime_audio:
599            return
600        packet = self._generate_packet_head(self._StopDumpAudioPage, 0)
601        self._audio_sock.send(packet)
602        # Drop audio pages until receive _StopDumpAudioPage response.
603        while True:
604            (message, _, _, _) = self._receive_whole_packet(self._audio_sock)
605            if message == self._get_response_type(self._StopDumpAudioPage):
606                break
607        self._is_realtime_audio = False
608