Improve logging.
[toast/stream2beamer.git] / lagarde.py
1 #!/usr/bin/python3
2
3 import argparse
4 import asyncio
5 import json
6 import logging
7 import ssl
8 import queue
9 from typing import Optional, List
10
11 import gi
12 import websockets
13
14 gi.require_version('Gst', '1.0')
15 from gi.repository import Gst
16
17 gi.require_version('GstWebRTC', '1.0')
18 from gi.repository import GstWebRTC
19
20 gi.require_version('GstSdp', '1.0')
21 from gi.repository import GstSdp
22
23 gi.require_version('GstRtspServer', '1.0')
24 from gi.repository import Gst, GstRtspServer, GObject, GLib
25
26 log = logging.getLogger(__name__)
27
28
29 class GstreamerRtspServer():
30     def __init__(self):
31         server = GstRtspServer.RTSPServer()
32         server.set_address("::")
33         server.set_service('8554')  # port as string
34         factory = GstRtspServer.RTSPMediaFactory()
35         # factory.set_launch("intervideosrc ! decodebin ! theoraenc ! queue ! rtptheorapay name=pay0")
36         # factory.set_launch("intervideosrc ! decodebin ! videoconvert ! video/x-raw,format=I420 ! theoraenc ! queue ! rtptheorapay name=pay0")
37         # factory.set_launch("intervideosrc ! decodebin ! videorate ! videoconvert ! video/x-raw,format=I420,framerate=10/1 ! theoraenc ! queue ! rtptheorapay name=pay0")
38         # factory.set_launch("intervideosrc ! decodebin ! videoconvert ! video/x-raw,format=I420 ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
39         factory.set_launch("intervideosrc ! decodebin ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
40         # factory.set_launch("intervideosrc ! decodebin ! videorate ! videoconvert ! video/x-raw,framerate=10/1 ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
41         factory.set_shared(True)
42         mountPoints = server.get_mount_points()
43         mountPoints.add_factory("/cug", factory)
44         server.attach()
45         self.server = server
46
47
48 class Lagarde:
49     def __init__(self):
50         self.sdp_offer: Optional[str] = None
51         self.websocket: Optional[websockets.client.WebSocketClientProtocol] = None
52         self.session_id = None
53         self.received_ice_candidates = queue.Queue()
54         self.generated_ice_candidates = queue.Queue()
55         self.user_fragments: Optional[List] = None
56         self.mids: Optional[List] = None
57         self.pipe = None
58         self.webrtcbin = None
59
60     def on_negotiation_needed(self, element):
61         log.info('on_negotiation_needed')
62
63     def on_ice_candidate(self, element, mlineindex, candidate):
64         log.info('on_ice_candidate')
65         self.generated_ice_candidates.put_nowait((mlineindex, candidate))
66
67     def webrtcbin_pad_added(self, element, pad):
68         log.info('webrtcbin_pad_added')
69         if pad.direction != Gst.PadDirection.SRC:
70             return
71         decodebin = Gst.ElementFactory.make('decodebin')
72         decodebin.connect('pad-added', self.decodebin_pad_added)
73         self.pipe.add(decodebin)
74         decodebin.sync_state_with_parent()
75         self.webrtcbin.link(decodebin)
76
77     def decodebin_pad_added(self, element, pad):
78         log.info('decodebin_pad_added')
79         if not pad.has_current_caps():
80             log.info(pad, 'has no caps, ignoring')
81             return
82
83         caps = pad.get_current_caps()
84         padsize = caps.get_size()
85         for i in range(padsize):
86             s = caps.get_structure(i)  # Gst.Structure
87             name = s.get_name()
88             if name.startswith('video'):
89                 q = Gst.ElementFactory.make('queue')
90                 conv = Gst.ElementFactory.make('videoconvert')
91                 sink = Gst.ElementFactory.make('intervideosink')
92                 self.pipe.add(q)
93                 self.pipe.add(conv)
94                 self.pipe.add(sink)
95                 self.pipe.sync_children_states()
96                 pad.link(q.get_static_pad('sink'))
97                 q.link(conv)
98                 conv.link(sink)
99                 self.pipe.set_state(Gst.State.PLAYING)
100
101     async def listen_to_gstreamer_bus(self):
102         self.webrtcbin = Gst.ElementFactory.make('webrtcbin', 'laplace')
103         self.pipe = Gst.Pipeline.new("pipeline")
104         Gst.Bin.do_add_element(self.pipe, self.webrtcbin)
105         bus = Gst.Pipeline.get_bus(self.pipe)
106         self.webrtcbin.connect('on-negotiation-needed', self.on_negotiation_needed)
107         self.webrtcbin.connect('on-ice-candidate', self.on_ice_candidate)
108         self.webrtcbin.connect('pad-added', self.webrtcbin_pad_added)
109         self.pipe.set_state(Gst.State.PLAYING)
110         try:
111             while True:
112                 if bus.have_pending():
113                     msg = bus.pop()
114                     if msg.type == Gst.MessageType.ERROR:
115                         log.error(f'Error from gstreamer message bus: {msg.get_structure()}')
116                         return
117                     elif msg.type == Gst.MessageType.EOS:  # end of stream
118                         log.info(f'Gstreamer message bus reports end of stream')
119                         return
120                 elif self.sdp_offer is not None:
121                     res, sm = GstSdp.SDPMessage.new()
122                     assert res == GstSdp.SDPResult.OK
123                     GstSdp.sdp_message_parse_buffer(bytes(self.sdp_offer.encode()), sm)
124                     # the three lines above can also be done this way in new versions of GStreamer:
125                     # sm = GstSdp.SDPMessage.new_from_text(sdp_offer)
126                     rd = GstWebRTC.WebRTCSessionDescription.new(GstWebRTC.WebRTCSDPType.OFFER, sm)
127                     gst_promise = Gst.Promise.new()
128                     self.webrtcbin.emit('set-remote-description', rd, gst_promise)
129                     await asyncio.get_event_loop().run_in_executor(None, gst_promise.wait)
130                     self.sdp_offer = None
131
132                     gst_promise = Gst.Promise.new()
133                     self.webrtcbin.emit('create-answer', None, gst_promise)
134                     result = await asyncio.get_event_loop().run_in_executor(None, gst_promise.wait)
135                     assert result == Gst.PromiseResult.REPLIED
136                     reply = gst_promise.get_reply()
137                     answer = reply.get_value('answer')
138                     sdp_message = answer.sdp
139                     self.user_fragments = [sdp_message.get_media(i).get_attribute_val('ice-ufrag')
140                                           for i in range(sdp_message.medias_len())]
141                     self.mids = [sdp_message.get_media(i).get_attribute_val('mid')
142                                           for i in range(sdp_message.medias_len())]
143                     sdp_answer = sdp_message.as_text()
144                     log.info(f'Send SDP answer')
145                     log.debug(f'SDP answer:\n{sdp_answer}')
146                     sdp_answer_msg = json.dumps({
147                         'SessionID': self.session_id,
148                         'Type': "gotAnswer",
149                         'Value': json.dumps({
150                             'type': 'answer',
151                             'sdp': sdp_answer
152                         })
153                     })
154                     gst_promise = Gst.Promise.new()
155                     self.webrtcbin.emit('set-local-description', answer, gst_promise)
156                     await asyncio.get_event_loop().run_in_executor(None, gst_promise.wait)
157                     gst_promise.get_reply()
158                     await self.websocket.send(sdp_answer_msg)
159
160                 elif self.received_ice_candidates.qsize() > 0:
161                     ic = self.received_ice_candidates.get_nowait()
162                     if ic['candidate'] != '':
163                         self.webrtcbin.emit('add-ice-candidate', ic['sdpMLineIndex'], ic['candidate'])
164
165                 elif self.generated_ice_candidates.qsize() > 0:
166                     mlineindex, candidate = self.generated_ice_candidates.get_nowait()
167                     icemsg_value = json.dumps({
168                         "candidate": candidate,
169                         "sdpMid": self.mids[mlineindex],
170                         "sdpMLineIndex": mlineindex,
171                         "usernameFragment": self.user_fragments[mlineindex],
172                     })
173                     icemsg = json.dumps({
174                         'SessionID': self.session_id,
175                         'Type': 'addCalleeIceCandidate',
176                         'Value': icemsg_value,
177                     })
178                     log.info(f'Send ICE candidate')
179                     log.debug(f'ICE candidate: {icemsg_value}')
180                     await self.websocket.send(icemsg)
181
182                 else:
183                     await asyncio.sleep(0.1)
184         finally:
185             self.pipe.set_state(Gst.State.NULL)
186
187     async def talk_to_websocket(self, uri, ssl_context):
188         async with websockets.connect(uri, ssl=ssl_context, close_timeout=0.5) as self.websocket:
189             async for msg in self.websocket:
190                 msg_json = json.loads(msg)
191                 msg_type = msg_json['Type']
192                 msg_value = msg_json['Value']
193                 assert self.session_id is None or self.session_id == msg_json['SessionID']
194                 if msg_type == 'newSession':
195                     self.session_id = msg_json['SessionID']
196                     log.info(f"New session {self.session_id}")
197                 elif msg_type == 'gotOffer':
198                     value_json = json.loads(msg_value)
199                     sdp = value_json['sdp']
200                     log.info(f'Got SDP offer')
201                     log.debug(f'SDP offer:\n{sdp}')
202                     self.sdp_offer = sdp
203                 elif msg_type == 'addCallerIceCandidate':
204                     value_json = json.loads(msg_value)
205                     log.info(f'Got ICE candidate')
206                     log.debug(f'ICE candidate: {value_json}')
207                     self.received_ice_candidates.put_nowait(value_json)
208                 elif msg_type == 'roomNotFound':
209                     log.error(f'The room was not found: {uri}')
210                     return
211                 elif msg_type == 'roomClosed':
212                     log.info(f'Oh noes, the room went away (session {self.session_id})!')
213                     self.session_id = None
214                     return
215                 else:
216                     log.error(f'Unknown message type {msg_type}')
217
218     async def talk_to_signaling_server(self, uri):
219         ssl_context = ssl.SSLContext()
220         ssl_context.check_hostname = False
221         ssl_context.verify_mode = ssl.CERT_NONE
222         while True:
223             await self.talk_to_websocket(uri, ssl_context)
224             await asyncio.sleep(0.1)
225
226     async def run(self, uri):
227         try:
228             talk_to_signaling_server_task = asyncio.Task(self.talk_to_signaling_server(uri))
229             listen_to_gstreamer_bus_task = asyncio.Task(self.listen_to_gstreamer_bus())
230             main_loop = asyncio.Task(gstreamer_main_loop())
231             done, pending = await asyncio.wait(
232                 [talk_to_signaling_server_task, listen_to_gstreamer_bus_task, main_loop],
233                 return_when=asyncio.FIRST_COMPLETED)
234             for d in done:
235                 d.result()
236             for p in pending:
237                 p.cancel()
238         except OSError as e:
239             print(e)
240
241
242 async def gstreamer_main_loop():
243     """Does the equivalent of the following lines in an async friendly way:
244         loop = GLib.MainLoop()
245         loop.run()
246     """
247     gst_loop = GLib.MainLoop()
248     context = gst_loop.get_context()
249     while True:
250         events_dispatched = context.iteration(False)
251         await asyncio.sleep(0. if events_dispatched else 0.01)
252
253
254 def main():
255     logging.basicConfig(level=logging.INFO, format='%(asctime)-15s %(message)s')
256     parser = argparse.ArgumentParser()
257     parser.add_argument('--uri', default='wss://localhost:1234/ws_connect?id=cug',
258                         help='Signalling server URI')
259     args = parser.parse_args()
260
261     Gst.init(None)
262     rtsp = GstreamerRtspServer()
263     lagarde = Lagarde()
264     asyncio.run(lagarde.run(args.uri), debug=True)
265
266
267 if __name__ == '__main__':
268     main()