ccbdbf1e616a7b1232bbf45943907d2e2b86d874
[toast/stream2beamer.git] / lagarde.py
1 #!/usr/bin/python3
2 import argparse
3 import asyncio
4 import json
5 import logging
6 import ssl
7 import queue
8 from typing import Optional, List
9
10 import gi
11 import websockets
12
13 gi.require_version('Gst', '1.0')
14 from gi.repository import Gst
15
16 gi.require_version('GstWebRTC', '1.0')
17 from gi.repository import GstWebRTC
18
19 gi.require_version('GstSdp', '1.0')
20 from gi.repository import GstSdp
21
22 gi.require_version('GstRtspServer', '1.0')
23 from gi.repository import Gst, GstRtspServer, GObject, GLib
24
25 log = logging.getLogger(__name__)
26
27
28 class RtspServer:
29     def __init__(self):
30         server = GstRtspServer.RTSPServer()
31         server.set_address("::")
32         server.set_service('8554')  # port as string
33         factory = GstRtspServer.RTSPMediaFactory()
34         # factory.set_launch("intervideosrc ! decodebin ! theoraenc ! queue ! rtptheorapay name=pay0")
35         # factory.set_launch("intervideosrc ! decodebin ! videoconvert ! video/x-raw,format=I420 ! theoraenc ! queue ! rtptheorapay name=pay0")
36         # factory.set_launch("intervideosrc ! decodebin ! videorate ! videoconvert ! video/x-raw,format=I420,framerate=10/1 ! theoraenc ! queue ! rtptheorapay name=pay0")
37         # factory.set_launch("intervideosrc ! decodebin ! videoconvert ! video/x-raw,format=I420 ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
38         factory.set_launch("intervideosrc ! decodebin ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
39         # factory.set_launch("intervideosrc ! decodebin ! videorate ! videoconvert ! video/x-raw,framerate=10/1 ! x264enc ! queue ! rtph264pay pt=96 name=pay0")
40         factory.set_shared(True)
41         mountPoints = server.get_mount_points()
42         mountPoints.add_factory("/cug", factory)
43         server.attach()
44         self.server = server
45
46
47 class Events:
48     def __init__(self):
49         self.sdp_offer = queue.Queue()
50         self.sdp_answer = queue.Queue()
51         self.generated_ice_candidates = queue.Queue()
52         self.received_ice_candidates = queue.Queue()
53         self.sdp_info = queue.Queue()  # (sdp_mids, user_fragments)
54         self.room_left = queue.Queue()
55
56
57 class SignalingClient:
58     def __init__(self, events: Events, uri):
59         self.events = events
60         self.uri = uri
61         self.ssl_context = ssl.SSLContext()
62         self.ssl_context.check_hostname = False
63         self.ssl_context.verify_mode = ssl.CERT_NONE
64         self.session_id = None
65
66     async def receive(self, uri):
67         async for msg in self.websocket:
68             msg_json = json.loads(msg)
69             msg_type = msg_json['Type']
70             msg_value = msg_json['Value']
71             assert self.session_id is None or self.session_id == msg_json['SessionID']
72             if msg_type == 'newSession':
73                 self.session_id = msg_json['SessionID']
74                 log.info(f"New session {self.session_id}")
75             elif msg_type == 'gotOffer':
76                 value_json = json.loads(msg_value)
77                 sdp = value_json['sdp']
78                 log.info(f'Got SDP offer')
79                 log.debug(f'SDP offer:\n{sdp}')
80                 self.events.sdp_offer.put_nowait(sdp)
81             elif msg_type == 'addCallerIceCandidate':
82                 value_json = json.loads(msg_value)
83                 log.info(f'Got ICE candidate')
84                 log.debug(f'ICE candidate: {value_json}')
85                 self.events.received_ice_candidates.put_nowait(value_json)
86             elif msg_type == 'roomNotFound':
87                 log.error(f'The room was not found: {uri}')
88                 return
89             elif msg_type == 'roomClosed':
90                 log.info(f'Oh noes, the room went away (session {self.session_id})!')
91                 self.events.room_left.put_nowait(True)
92                 return
93             else:
94                 log.error(f'Unknown message type {msg_type}')
95
96     async def send(self):
97         sdp_mids = None
98         user_fragments = None
99         while True:
100             if self.events.sdp_answer.qsize() > 0:
101                 sdp_answer = self.events.sdp_answer.get_nowait()
102                 sdp_answer_msg = json.dumps({
103                     'SessionID': self.session_id,
104                     'Type': "gotAnswer",
105                     'Value': json.dumps({
106                         'type': 'answer',
107                         'sdp': sdp_answer
108                     })
109                 })
110                 await self.websocket.send(sdp_answer_msg)
111
112             elif self.events.sdp_info.qsize() > 0:
113                 sdp_mids, user_fragments = self.events.sdp_info.get_nowait()
114
115             elif self.events.generated_ice_candidates.qsize() > 0 \
116                     and sdp_mids is not None and user_fragments is not None:
117                 mlineindex, candidate = self.events.generated_ice_candidates.get_nowait()
118                 sdp_mid = sdp_mids[mlineindex]
119                 user_fragment = user_fragments[mlineindex]
120                 icemsg_value = json.dumps({
121                     "candidate": candidate,
122                     "sdpMid": sdp_mid,
123                     "sdpMLineIndex": mlineindex,
124                     "usernameFragment": user_fragment,
125                 })
126                 icemsg = json.dumps({
127                     'SessionID': self.session_id,
128                     'Type': 'addCalleeIceCandidate',
129                     'Value': icemsg_value,
130                 })
131                 log.info(f'Send ICE candidate')
132                 log.debug(f'ICE candidate: {icemsg_value}')
133                 await self.websocket.send(icemsg)
134
135             else:
136                 await asyncio.sleep(0.2)
137
138     async def run(self):
139         self.session_id = None
140         async with websockets.connect(self.uri, ssl=self.ssl_context, close_timeout=0.5) as self.websocket:
141             receive_task = asyncio.Task(self.receive(self.uri))
142             send_task = asyncio.Task(self.send())
143             done, pending = await asyncio.wait([receive_task, send_task], return_when=asyncio.FIRST_COMPLETED)
144             for task in pending:
145                 task.cancel()
146
147
148 class WebRTCClient:
149     def __init__(self, events: Events):
150         self.events = events
151         self.webrtcbin = Gst.ElementFactory.make('webrtcbin', 'laplace')
152         self.pipe = Gst.Pipeline.new("pipeline")
153         Gst.Bin.do_add_element(self.pipe, self.webrtcbin)
154         self.webrtcbin.connect('on-negotiation-needed', self.on_negotiation_needed)
155         self.webrtcbin.connect('on-ice-candidate', self.on_ice_candidate)
156         self.webrtcbin.connect('pad-added', self.webrtcbin_pad_added)
157
158     def on_negotiation_needed(self, element):
159         log.info('on_negotiation_needed')
160
161     def on_ice_candidate(self, element, mlineindex, candidate):
162         log.info('on_ice_candidate')
163         self.events.generated_ice_candidates.put_nowait((mlineindex, candidate))
164
165     def webrtcbin_pad_added(self, element, pad):
166         log.info('webrtcbin_pad_added')
167         if pad.direction != Gst.PadDirection.SRC:
168             return
169         decodebin = Gst.ElementFactory.make('decodebin')
170         decodebin.connect('pad-added', self.decodebin_pad_added)
171         self.pipe.add(decodebin)
172         decodebin.sync_state_with_parent()
173         self.webrtcbin.link(decodebin)
174
175     def decodebin_pad_added(self, element, pad):
176         log.info('decodebin_pad_added')
177         if not pad.has_current_caps():
178             log.info(pad, 'has no caps, ignoring')
179             return
180         caps = pad.get_current_caps()
181         padsize = caps.get_size()
182
183         log.info(f'>>>> {padsize} {caps}')
184
185         for i in range(padsize):
186             s = caps.get_structure(i)  # Gst.Structure
187             name = s.get_name()
188             log.info(f'###### {name}')
189             if name.startswith('video'):
190                 q = Gst.ElementFactory.make('queue')
191                 conv = Gst.ElementFactory.make('videoconvert')
192                 enc = Gst.ElementFactory.make('x264enc')
193                 enc.set_property('bitrate', 1000)
194                 enc.set_property('tune', 'zerolatency')
195                 capsfilter = Gst.ElementFactory.make('capsfilter')
196                 capsfilter.set_properties(Gst.Caps.from_string('video/x-h264,stream-format=(string)avc'))
197                 flmux = Gst.ElementFactory.make('flvmux')
198                 sink = Gst.ElementFactory.make('rtmpsink')
199                 sink.set_property('location', 'rtmp://192.168.1.46:1935/gregoa')
200                 # sink.set_property('location', 'rtmp://bla:1936/gregoa')
201                 print(sink.props.location, dir(sink.props))
202                 assert q and conv and enc and capsfilter and flmux and sink
203
204                 self.pipe.add(q)
205                 self.pipe.add(conv)
206                 self.pipe.add(enc)
207                 self.pipe.add(capsfilter)
208                 self.pipe.add(flmux)
209                 self.pipe.add(sink)
210                 self.pipe.sync_children_states()
211
212                 q_pad_sink = q.get_static_pad('sink')
213                 assert q_pad_sink
214                 pad_link_return = pad.link(q_pad_sink)
215                 assert pad_link_return == Gst.PadLinkReturn.OK
216
217                 # ok = element.link(q)
218                 # assert ok
219
220                 ok = q.link(conv)
221                 assert ok
222                 ok = conv.link(enc)
223                 assert ok
224                 ok = enc.link(capsfilter)
225                 assert ok
226                 ok = capsfilter.link(flmux)
227                 assert ok
228                 ok = flmux.link(sink)
229                 assert ok
230                 self.pipe.set_state(Gst.State.PLAYING)
231                 #print(dir(Gst.DebugGraphDetails))
232                 #Gst.debug_bin_to_dot_data(element, Gst.DebugGraphDetails.ALL)
233
234             elif name.startswith('audio'):
235                 q = Gst.ElementFactory.make('queue')
236                 conv = Gst.ElementFactory.make('audioconvert')
237                 resample = Gst.ElementFactory.make('audioresample')
238                 sink = Gst.ElementFactory.make('autoaudiosink')
239                 self.pipe.add(q)
240                 self.pipe.add(conv)
241                 self.pipe.add(resample)
242                 self.pipe.add(sink)
243                 self.pipe.sync_children_states()
244                 pad.link(q.get_static_pad('sink'))
245                 q.link(conv)
246                 conv.link(resample)
247                 resample.link(sink)
248
249     def set_remote_desciption_done(self, gst_promise):
250         gst_promise = Gst.Promise.new_with_change_func(self.create_answer_done)
251         self.webrtcbin.emit('create-answer', None, gst_promise)
252
253     def create_answer_done(self, gst_promise):
254         reply = gst_promise.get_reply()
255         answer = reply.get_value('answer')
256         sdp_message = answer.sdp
257         mids = [sdp_message.get_media(i).get_attribute_val('mid')
258                      for i in range(sdp_message.medias_len())]
259         user_fragments = [sdp_message.get_media(i).get_attribute_val('ice-ufrag')
260                                for i in range(sdp_message.medias_len())]
261         self.events.sdp_info.put_nowait((mids, user_fragments))
262         sdp_answer = sdp_message.as_text()
263         log.info(f'Send SDP answer')
264         log.debug(f'SDP answer:\n{sdp_answer}')
265         self.events.sdp_answer.put_nowait(sdp_answer)
266         gst_promise = Gst.Promise.new_with_change_func(self.set_local_description_done)
267         self.webrtcbin.emit('set-local-description', answer, gst_promise)
268
269     def set_local_description_done(self, gst_promise):
270         gst_promise.get_reply()
271
272     async def run(self):
273         bus = Gst.Pipeline.get_bus(self.pipe)
274         self.pipe.set_state(Gst.State.PLAYING)
275         try:
276             while True:
277                 if bus.have_pending():
278                     msg = bus.pop()
279                     if msg.type == Gst.MessageType.ERROR:
280                         log.error(f'Error from gstreamer message bus: {msg.get_structure()}')
281                         return
282                     elif msg.type == Gst.MessageType.EOS:  # end of stream
283                         log.info(f'Gstreamer message bus reports end of stream')
284                         return
285                 elif self.events.sdp_offer.qsize() > 0:
286                     sdp_offer = self.events.sdp_offer.get_nowait()
287                     res, sm = GstSdp.SDPMessage.new()
288                     assert res == GstSdp.SDPResult.OK
289                     GstSdp.sdp_message_parse_buffer(bytes(sdp_offer.encode()), sm)
290                     # the three lines above can also be done this way in new versions of GStreamer:
291                     # sm = GstSdp.SDPMessage.new_from_text(sdp_offer)
292                     rd = GstWebRTC.WebRTCSessionDescription.new(GstWebRTC.WebRTCSDPType.OFFER, sm)
293                     gst_promise = Gst.Promise.new_with_change_func(self.set_remote_desciption_done)
294                     self.webrtcbin.emit('set-remote-description', rd, gst_promise)
295
296                 elif self.events.received_ice_candidates.qsize() > 0:
297                     ic = self.events.received_ice_candidates.get_nowait()
298                     if ic['candidate'] != '':
299                         self.webrtcbin.emit('add-ice-candidate', ic['sdpMLineIndex'], ic['candidate'])
300
301                 elif self.events.room_left.qsize() > 0:
302                     self.events.room_left.get_nowait()
303                     return
304
305                 else:
306                     await asyncio.sleep(0.1)
307         finally:
308             self.pipe.set_state(Gst.State.NULL)
309
310
311 async def run_repeated(task):
312     while True:
313         await task()
314         await asyncio.sleep(0.1)
315
316
317 async def run(uri):
318     try:
319         events = Events()
320         # rtsp = RtspServer()
321         webrtc = WebRTCClient(events)
322         signaling = SignalingClient(events, uri)
323
324         webrtc_task = asyncio.Task(webrtc.run())
325         signaling_task = asyncio.Task(signaling.run())
326
327         done, pending = await asyncio.wait([webrtc_task, signaling_task],
328             return_when=asyncio.FIRST_COMPLETED)
329
330         for task in done:
331             task.result()
332         for task in pending:
333             task.cancel()
334     except OSError as e:
335         print(e)
336
337
338 def main():
339     logging.basicConfig(level=logging.DEBUG, format='%(asctime)-15s %(message)s')
340     parser = argparse.ArgumentParser()
341     parser.add_argument('--uri', default='wss://localhost:1234/ws_connect?id=cug',
342                         help='Signalling server URI')
343     args = parser.parse_args()
344
345     Gst.init(None)
346     asyncio.run(run(args.uri), debug=True)
347
348
349 if __name__ == '__main__':
350     main()