Yes, it's as pythonic as such an non-Python object can be.
[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 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 log = logging.getLogger(__name__)
23
24
25 class Lagarde:
26     def __init__(self):
27         self.sdp_offer: Optional[str] = None
28         self.websocket: Optional[websockets.client.WebSocketClientProtocol] = None
29         self.session_id = None
30         self.received_ice_candidates = []
31         self.generated_ice_candidates = []
32         self.user_fragments: Optional[List] = None
33         self.mids: Optional[List] = None
34         self.pipe = None
35         self.webrtcbin = None
36
37     def on_negotiation_needed(self, element):
38         log.info('on_negotiation_needed')
39
40     def on_ice_candidate(self, element, mlineindex, candidate):
41         log.info('on_ice_candidate')
42         self.generated_ice_candidates.append((mlineindex, candidate))
43
44     def webrtcbin_pad_added(self, element, pad):
45         log.info('webrtcbin_pad_added')
46         if pad.direction != Gst.PadDirection.SRC:
47             return
48         decodebin = Gst.ElementFactory.make('decodebin')
49         decodebin.connect('pad-added', self.decodebin_pad_added)
50         self.pipe.add(decodebin)
51         decodebin.sync_state_with_parent()
52         self.webrtcbin.link(decodebin)
53
54     def decodebin_pad_added(self, element, pad):
55         log.info('decodebin_pad_added')
56         if not pad.has_current_caps():
57             log.info(pad, 'has no caps, ignoring')
58             return
59
60         caps = pad.get_current_caps()
61         padsize = caps.get_size()
62         for i in range(padsize):
63             s = caps.get_structure(i) # Gst.Structure
64             name = s.get_name()
65             if name.startswith('video'):
66                 q = Gst.ElementFactory.make('queue')
67                 conv = Gst.ElementFactory.make('videoconvert')
68                 # sink = Gst.ElementFactory.make('autovideosink') # needs XDG_RUNTIME_DIR
69                 sink = Gst.ElementFactory.make('xvimagesink')
70                 self.pipe.add(q)
71                 self.pipe.add(conv)
72                 self.pipe.add(sink)
73                 self.pipe.sync_children_states()
74                 pad.link(q.get_static_pad('sink'))
75                 q.link(conv)
76                 conv.link(sink)
77             elif name.startswith('audio'):
78                 q = Gst.ElementFactory.make('queue')
79                 conv = Gst.ElementFactory.make('audioconvert')
80                 resample = Gst.ElementFactory.make('audioresample')
81                 sink = Gst.ElementFactory.make('autoaudiosink')
82                 self.pipe.add(q)
83                 self.pipe.add(conv)
84                 self.pipe.add(resample)
85                 self.pipe.add(sink)
86                 self.pipe.sync_children_states()
87                 pad.link(q.get_static_pad('sink'))
88                 q.link(conv)
89                 conv.link(resample)
90                 resample.link(sink)
91
92     async def listen_to_gstreamer_bus(self):
93         Gst.init(None)
94         self.webrtcbin = Gst.ElementFactory.make('webrtcbin', 'laplace')
95         self.pipe = Gst.Pipeline.new("pipeline")
96         Gst.Bin.do_add_element(self.pipe, self.webrtcbin)
97         bus = Gst.Pipeline.get_bus(self.pipe)
98         self.webrtcbin.connect('on-negotiation-needed', self.on_negotiation_needed)
99         self.webrtcbin.connect('on-ice-candidate', self.on_ice_candidate)
100         self.webrtcbin.connect('pad-added', self.webrtcbin_pad_added)
101         self.pipe.set_state(Gst.State.PLAYING)
102         try:
103             while True:
104                 if bus.have_pending():
105                     msg = bus.pop()  # Gst.Message, has to be unref'ed.
106                     if msg.type != Gst.MessageType.STATE_CHANGED:
107                         # log.info(f'Receive Gst.Message: {msg.type}, {msg.seqnum}, {msg.get_structure()}')
108                         # log.info(f'{webrtcbin.props.signaling_state} {webrtcbin.props.ice_gathering_state} {webrtcbin.props.ice_connection_state}')
109                         # Gst.Message.unref(msg)
110                         pass
111                 elif self.sdp_offer is not None:
112                     res, sm = GstSdp.SDPMessage.new()
113                     assert res == GstSdp.SDPResult.OK
114                     GstSdp.sdp_message_parse_buffer(bytes(self.sdp_offer.encode()), sm)
115                     # the three lines above can also be done this way in new versions of GStreamer:
116                     # sm = GstSdp.SDPMessage.new_from_text(sdp_offer)
117                     rd = GstWebRTC.WebRTCSessionDescription.new(GstWebRTC.WebRTCSDPType.OFFER, sm)
118                     gst_promise = Gst.Promise.new()
119                     self.webrtcbin.emit('set-remote-description', rd, gst_promise)
120                     gst_promise.wait()
121                     self.sdp_offer = None
122
123                     log.info('create-answer')
124                     gst_promise = Gst.Promise.new()
125                     self.webrtcbin.emit('create-answer', None, gst_promise)
126                     result = gst_promise.wait()
127                     assert result == Gst.PromiseResult.REPLIED
128                     reply = gst_promise.get_reply()
129                     answer = reply.get_value('answer')
130                     sdp_message = answer.sdp
131                     self.user_fragments = [sdp_message.get_media(i).get_attribute_val('ice-ufrag')
132                                           for i in range(sdp_message.medias_len())]
133                     self.mids = [sdp_message.get_media(i).get_attribute_val('mid')
134                                           for i in range(sdp_message.medias_len())]
135                     sdp_answer = sdp_message.as_text()
136                     log.info(sdp_answer)
137                     sdp_answer_msg = json.dumps({
138                         'SessionID': self.session_id,
139                         'Type': "gotAnswer",
140                         'Value': json.dumps({
141                             'type': 'answer',
142                             'sdp': sdp_answer
143                         })
144                     })
145                     gst_promise = Gst.Promise.new()
146                     self.webrtcbin.emit('set-local-description', answer, gst_promise)
147                     gst_promise.wait()
148                     gst_promise.get_reply()
149                     await self.websocket.send(sdp_answer_msg)
150
151                 elif len(self.received_ice_candidates) > 0:
152                     ic = self.received_ice_candidates.pop(0)
153                     if ic['candidate'] != '':
154                         self.webrtcbin.emit('add-ice-candidate', ic['sdpMLineIndex'], ic['candidate'])
155
156                 elif len(self.generated_ice_candidates) > 0:
157                     mlineindex, candidate = self.generated_ice_candidates.pop(0)
158                     icemsg = json.dumps({
159                         'SessionID': self.session_id,
160                         'Type': 'addCalleeIceCandidate',
161                         'Value': json.dumps({
162                             "candidate": candidate,
163                             "sdpMid": self.mids[mlineindex],
164                             "sdpMLineIndex": mlineindex,
165                             "usernameFragment": self.user_fragments[mlineindex],
166                         })
167                     })
168                     log.info(f'send_ice_candidate_message with {icemsg}')
169                     await self.websocket.send(icemsg)
170
171                 else:
172                     await asyncio.sleep(0.1)
173         finally:
174             self.pipe.set_state(Gst.State.NULL)
175
176     async def talk_to_websocket(self, uri):
177         ssl_context = ssl.SSLContext()
178         ssl_context.check_hostname = False
179         ssl_context.verify_mode = ssl.CERT_NONE
180         async with websockets.connect(uri, ssl=ssl_context) as self.websocket:
181             async for msg in self.websocket:
182                 msg_json = json.loads(msg)
183                 msg_type = msg_json['Type']
184                 msg_value = msg_json['Value']
185                 self.session_id = msg_json['SessionID']
186                 log.info(f"receive for session {self.session_id} type {msg_type}")
187                 if msg_type == 'newSession':
188                     pass
189                 elif msg_type == 'gotOffer':
190                     value_json = json.loads(msg_value)
191                     sdp = value_json['sdp']
192                     log.info(f'SDP: {sdp}')
193                     self.sdp_offer = sdp
194                 elif msg_type == 'addCallerIceCandidate':
195                     value_json = json.loads(msg_value)
196                     log.info(f'ICE: {value_json}')
197                     self.received_ice_candidates.append(value_json)
198                 elif msg_type == 'roomClosed':
199                     log.info('Oh noes, the room went away!')
200                     # and here we should clean up
201                 else:
202                     log.error(f'Unknown message type {msg_type}')
203
204     async def run(self, uri):
205         talk_to_websocket_task = asyncio.Task(self.talk_to_websocket(uri))
206         listen_to_gstreamer_bus_task = asyncio.Task(self.listen_to_gstreamer_bus())
207         done, pending = await asyncio.wait(
208             [talk_to_websocket_task, listen_to_gstreamer_bus_task],
209             return_when=asyncio.FIRST_COMPLETED)
210         for d in done:
211             d.result()
212         for p in pending:
213             p.cancel()
214
215
216 def main():
217     logging.basicConfig(level=logging.INFO, format='%(asctime)-15s %(message)s')
218     parser = argparse.ArgumentParser()
219     parser.add_argument('--uri', default='wss://localhost:1234/ws_connect?id=cug',
220                         help='Signalling server URI')
221     args = parser.parse_args()
222     lagarde = Lagarde()
223     asyncio.run(lagarde.run(args.uri), debug=True)
224
225
226 if __name__ == '__main__':
227     main()