Coverage for /builds/BuildGrid/buildgrid/buildgrid/server/utils/digests.py: 95.83%

1# Copyright (C) 2024 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 

16from dataclasses import dataclass 

17from typing import Optional 

18 

19from buildgrid._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Digest, DigestFunction 

20from buildgrid.server.settings import HASH, HASH_LENGTH 

21 

22 

23@dataclass(frozen=True) 

24class HashableDigest: 

25 hash: str 

26 size_bytes: int 

27 

28 def to_digest(self) -> Digest: 

29 return Digest(hash=self.hash, size_bytes=self.size_bytes) 

30 

31 

32def get_hash_type() -> "DigestFunction.Value.ValueType": 

33 """Returns the hash type.""" 

34 hash_name = HASH().name 

35 if hash_name == "sha256": 

36 return DigestFunction.SHA256 

37 return DigestFunction.UNKNOWN 

38 

39 

40def create_digest(bytes_to_digest: bytes) -> Digest: 

41 """Computes the :obj:`Digest` of a piece of data. 

42 

43 The :obj:`Digest` of a data is a function of its hash **and** size. 

44 

45 Args: 

46 bytes_to_digest (bytes): byte data to digest. 

47 

48 Returns: 

49 :obj:`Digest`: The :obj:`Digest` for the given byte data. 

50 """ 

51 return Digest(hash=HASH(bytes_to_digest).hexdigest(), size_bytes=len(bytes_to_digest)) 

52 

53 

54def parse_digest(digest_string: str) -> Optional[Digest]: 

55 """Creates a :obj:`Digest` from a digest string. 

56 

57 A digest string should alway be: ``{hash}/{size_bytes}``. 

58 

59 Args: 

60 digest_string (str): the digest string. 

61 

62 Returns: 

63 :obj:`Digest`: The :obj:`Digest` read from the string or None if 

64 `digest_string` is not a valid digest string. 

65 """ 

66 digest_hash, digest_size = digest_string.split("/") 

67 

68 if len(digest_hash) == HASH_LENGTH and digest_size.isdigit(): 

69 return Digest(hash=digest_hash, size_bytes=int(digest_size)) 

70 

71 return None 

72 

73 

74def validate_digest_data(digest: Digest, data: bytes) -> bool: 

75 """Validate that the given digest corresponds to the given data.""" 

76 return len(data) == digest.size_bytes and HASH(data).hexdigest() == digest.hash