Coverage for /builds/BuildGrid/buildgrid/buildgrid/client/authentication.py: 67.57%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

74 statements  

1# Copyright (C) 2018 Bloomberg LP 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# <http://www.apache.org/licenses/LICENSE-2.0> 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15 

16import base64 

17from collections import namedtuple 

18import os 

19from typing import Optional 

20 

21import grpc 

22from grpc import aio # type: ignore 

23 

24from buildgrid._exceptions import InvalidArgumentError 

25from buildgrid.utils import read_file 

26 

27 

28def load_tls_channel_credentials(client_key=None, client_cert=None, server_cert=None): 

29 """Looks-up and loads TLS gRPC client channel credentials. 

30 

31 Args: 

32 client_key(str, optional): Client certificate chain file path. 

33 client_cert(str, optional): Client private key file path. 

34 server_cert(str, optional): Serve root certificate file path. 

35 

36 Returns: 

37 ChannelCredentials: Credentials to be used for a TLS-encrypted gRPC 

38 client channel. 

39 """ 

40 if server_cert and os.path.exists(server_cert): 

41 server_cert_pem = read_file(server_cert) 

42 else: 

43 server_cert_pem = None 

44 server_cert = None 

45 

46 if client_key and os.path.exists(client_key): 

47 client_key_pem = read_file(client_key) 

48 else: 

49 client_key_pem = None 

50 client_key = None 

51 

52 if client_key_pem and client_cert and os.path.exists(client_cert): 

53 client_cert_pem = read_file(client_cert) 

54 else: 

55 client_cert_pem = None 

56 client_cert = None 

57 

58 credentials = grpc.ssl_channel_credentials(root_certificates=server_cert_pem, 

59 private_key=client_key_pem, 

60 certificate_chain=client_cert_pem) 

61 

62 return credentials, (client_key, client_cert, server_cert,) 

63 

64 

65def load_channel_authorization_token(auth_token=None): 

66 """Looks-up and loads client authorization token. 

67 

68 Args: 

69 auth_token (str, optional): Token file path. 

70 

71 Returns: 

72 str: Encoded token string. 

73 """ 

74 if auth_token and os.path.exists(auth_token): 

75 return read_file(auth_token).decode() 

76 

77 # TODO: Try loading the token from a default location? 

78 

79 return None 

80 

81 

82class AuthMetadataClientInterceptorBase: 

83 

84 def __init__(self, auth_token: Optional[str]=None, auth_secret: Optional[bytes]=None): 

85 """Initialises a new :class:`AuthMetadataClientInterceptorBase`. 

86 

87 Important: 

88 One of `auth_token` or `auth_secret` must be provided. 

89 

90 Args: 

91 auth_token (str, optional): Authorization token as a string. 

92 auth_secret (bytes, optional): Authorization secret as bytes. 

93 

94 Raises: 

95 InvalidArgumentError: If neither `auth_token` or `auth_secret` are 

96 provided. 

97 """ 

98 if auth_token: 

99 self.__secret = auth_token.strip() 

100 

101 elif auth_secret: 

102 self.__secret = base64.b64encode(auth_secret.strip()).decode() 

103 

104 else: 

105 raise InvalidArgumentError("A secret or token must be provided") 

106 

107 self.__header_field_name = 'authorization' 

108 self.__header_field_value = f'Bearer {self.__secret}' 

109 

110 def _amend_call_details(self, client_call_details, grpc_call_details_class): 

111 """Appends an authorization field to given client call details.""" 

112 if client_call_details.metadata is not None: 

113 new_metadata = list(client_call_details.metadata) 

114 else: 

115 new_metadata = [] 

116 

117 new_metadata.append((self.__header_field_name, self.__header_field_value,)) 

118 

119 class _ClientCallDetails( 

120 namedtuple('_ClientCallDetails', 

121 ('method', 'timeout', 'credentials', 'metadata', 'wait_for_ready',)), 

122 grpc_call_details_class): 

123 pass 

124 

125 return _ClientCallDetails(client_call_details.method, 

126 client_call_details.timeout, 

127 client_call_details.credentials, 

128 new_metadata, 

129 client_call_details.wait_for_ready) 

130 

131 

132class AuthMetadataClientInterceptor(AuthMetadataClientInterceptorBase, 

133 grpc.UnaryUnaryClientInterceptor, 

134 grpc.UnaryStreamClientInterceptor, 

135 grpc.StreamUnaryClientInterceptor, 

136 grpc.StreamStreamClientInterceptor): 

137 

138 def __init__(self, auth_token: Optional[str]=None, auth_secret: Optional[bytes]=None): 

139 AuthMetadataClientInterceptorBase.__init__(self, auth_token, auth_secret) 

140 

141 def intercept_unary_unary(self, continuation, client_call_details, request): 

142 new_details = self._amend_call_details(client_call_details, grpc.ClientCallDetails) 

143 

144 return continuation(new_details, request) 

145 

146 def intercept_unary_stream(self, continuation, client_call_details, request): 

147 new_details = self._amend_call_details(client_call_details, grpc.ClientCallDetails) 

148 

149 return continuation(new_details, request) 

150 

151 def intercept_stream_unary(self, continuation, client_call_details, request_iterator): 

152 new_details = self._amend_call_details(client_call_details, grpc.ClientCallDetails) 

153 

154 return continuation(new_details, request_iterator) 

155 

156 def intercept_stream_stream(self, continuation, client_call_details, request_iterator): 

157 new_details = self._amend_call_details(client_call_details, grpc.ClientCallDetails) 

158 

159 return continuation(new_details, request_iterator) 

160 

161 

162class AsyncAuthMetadataClientInterceptor(AuthMetadataClientInterceptorBase, 

163 aio.UnaryUnaryClientInterceptor, 

164 aio.UnaryStreamClientInterceptor, 

165 aio.StreamUnaryClientInterceptor, 

166 aio.StreamStreamClientInterceptor): 

167 

168 def __init__(self, auth_token: Optional[str]=None, auth_secret: Optional[bytes]=None): 

169 AuthMetadataClientInterceptorBase.__init__(self, auth_token, auth_secret) 

170 

171 async def intercept_unary_unary(self, continuation, client_call_details, request): 

172 new_details = self._amend_call_details(client_call_details, aio.ClientCallDetails) 

173 

174 return await continuation(new_details, request) 

175 

176 async def intercept_unary_stream(self, continuation, client_call_details, request): 

177 new_details = self._amend_call_details(client_call_details, aio.ClientCallDetails) 

178 

179 return await continuation(new_details, request) 

180 

181 async def intercept_stream_unary(self, continuation, client_call_details, request_iterator): 

182 new_details = self._amend_call_details(client_call_details, aio.ClientCallDetails) 

183 

184 return await continuation(new_details, request_iterator) 

185 

186 async def intercept_stream_stream(self, continuation, client_call_details, request_iterator): 

187 new_details = self._amend_call_details(client_call_details, aio.ClientCallDetails) 

188 

189 return await continuation(new_details, request_iterator)