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

1import logging 

2 

3import grpc 

4from google.protobuf import empty_pb2 

5from sqlalchemy import select 

6from sqlalchemy.orm import Session 

7 

8from couchers.context import CouchersContext 

9from couchers.crypto import secure_compare 

10from couchers.interceptors import MediaInterceptor 

11from couchers.models import InitiatedUpload, Upload 

12from couchers.proto import media_pb2, media_pb2_grpc 

13 

14logger = logging.getLogger(__name__) 

15 

16 

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

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

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

20 

21 return MediaInterceptor(is_authorized) 

22 

23 

24class Media(media_pb2_grpc.MediaServicer): 

25 def UploadConfirmation( 

26 self, request: media_pb2.UploadConfirmationReq, context: CouchersContext, session: Session 

27 ) -> empty_pb2.Empty: 

28 initiated_upload = session.execute( 

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

30 ).scalar_one_or_none() 

31 

32 if not initiated_upload: 

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

34 

35 # move it to a completed upload 

36 upload = Upload( 

37 key=request.key, 

38 filename=request.filename, 

39 creator_user_id=initiated_upload.initiator_user_id, 

40 ) 

41 session.add(upload) 

42 

43 # delete the old upload 

44 session.delete(initiated_upload) 

45 

46 return empty_pb2.Empty()