Coverage for src/couchers/servicers/media.py: 95%

1import logging 

2 

3import grpc 

4from google.protobuf import empty_pb2 

5 

6from couchers.crypto import secure_compare 

7from couchers.interceptors import MediaInterceptor 

8from couchers.models import InitiatedUpload, Upload 

9from couchers.proto import media_pb2_grpc 

10from couchers.sql import couchers_select as select 

11 

12logger = logging.getLogger(__name__) 

13 

14 

15def get_media_auth_interceptor(secret_token: str) -> MediaInterceptor: 

16 def is_authorized(token: str) -> bool: 

17 return secure_compare(token.encode("ascii"), secret_token.encode("ascii")) 

18 

19 return MediaInterceptor(is_authorized) 

20 

21 

22class Media(media_pb2_grpc.MediaServicer): 

23 def UploadConfirmation(self, request, context, session) -> empty_pb2.Empty: 

24 initiated_upload = session.execute( 

25 select(InitiatedUpload).where(InitiatedUpload.key == request.key).where(InitiatedUpload.is_valid) 

26 ).scalar_one_or_none() 

27 

28 if not initiated_upload: 

29 context.abort_with_error_code(grpc.StatusCode.NOT_FOUND, "upload_not_found") 

30 

31 # move it to a completed upload 

32 upload = Upload( 

33 key=request.key, 

34 filename=request.filename, 

35 creator_user_id=initiated_upload.initiator_user_id, 

36 ) 

37 session.add(upload) 

38 

39 # delete the old upload 

40 session.delete(initiated_upload) 

41 

42 return empty_pb2.Empty()