1616from azure .core .exceptions import HttpResponseError , ServiceResponseError
1717
1818from azure .core .tracing .common import with_current_context
19- from ._shared .encryption import decrypt_blob
19+ from ._shared .encryption import (
20+ adjust_blob_size_for_encryption ,
21+ decrypt_blob ,
22+ get_adjusted_download_range_and_offset ,
23+ is_encryption_v2 ,
24+ parse_encryption_data
25+ )
2026from ._shared .request_handlers import validate_and_format_range_headers
2127from ._shared .response_handlers import process_storage_error , parse_length_from_content_range
22- from ._deserialize import get_page_ranges_result
28+ from ._deserialize import deserialize_blob_properties , get_page_ranges_result
2329
2430
25- def process_range_and_offset (start_range , end_range , length , encryption ):
31+ def process_range_and_offset (start_range , end_range , length , encryption_options , encryption_data ):
2632 start_offset , end_offset = 0 , 0
27- if encryption .get ("key" ) is not None or encryption .get ("resolver" ) is not None :
28- if start_range is not None :
29- # Align the start of the range along a 16 byte block
30- start_offset = start_range % 16
31- start_range -= start_offset
32-
33- # Include an extra 16 bytes for the IV if necessary
34- # Because of the previous offsetting, start_range will always
35- # be a multiple of 16.
36- if start_range > 0 :
37- start_offset += 16
38- start_range -= 16
39-
40- if length is not None :
41- # Align the end of the range along a 16 byte block
42- end_offset = 15 - (end_range % 16 )
43- end_range += end_offset
33+ if encryption_options .get ("key" ) is not None or encryption_options .get ("resolver" ) is not None :
34+ return get_adjusted_download_range_and_offset (
35+ start_range ,
36+ end_range ,
37+ length ,
38+ encryption_data )
4439
4540 return (start_range , end_range ), (start_offset , end_offset )
4641
@@ -81,6 +76,7 @@ def __init__(
8176 parallel = None ,
8277 validate_content = None ,
8378 encryption_options = None ,
79+ encryption_data = None ,
8480 progress_hook = None ,
8581 ** kwargs
8682 ):
@@ -108,6 +104,7 @@ def __init__(
108104
109105 # Encryption
110106 self .encryption_options = encryption_options
107+ self .encryption_data = encryption_data
111108
112109 # Parameters for each get operation
113110 self .validate_content = validate_content
@@ -183,7 +180,7 @@ def _do_optimize(self, given_range_start, given_range_end):
183180
184181 def _download_chunk (self , chunk_start , chunk_end ):
185182 download_range , offset = process_range_and_offset (
186- chunk_start , chunk_end , chunk_end , self .encryption_options
183+ chunk_start , chunk_end , chunk_end , self .encryption_options , self . encryption_data
187184 )
188185
189186 # No need to download the empty chunk from server if there's no data in the chunk to be downloaded.
@@ -335,6 +332,10 @@ def __init__(
335332 self ._file_size = None
336333 self ._non_empty_ranges = None
337334 self ._response = None
335+ self ._encryption_data = None
336+
337+ if self ._encryption_options .get ("key" ) is not None or self ._encryption_options .get ("resolver" ) is not None :
338+ self ._get_encryption_data_request ()
338339
339340 # The service only provides transactional MD5s for chunks under 4MB.
340341 # If validate_content is on, get only self.MAX_CHUNK_GET_SIZE for the first
@@ -349,7 +350,11 @@ def __init__(
349350 initial_request_end = initial_request_start + self ._first_get_size - 1
350351
351352 self ._initial_range , self ._initial_offset = process_range_and_offset (
352- initial_request_start , initial_request_end , self ._end_range , self ._encryption_options
353+ initial_request_start ,
354+ initial_request_end ,
355+ self ._end_range ,
356+ self ._encryption_options ,
357+ self ._encryption_data
353358 )
354359
355360 self ._response = self ._initial_request ()
@@ -376,6 +381,21 @@ def __init__(
376381 def __len__ (self ):
377382 return self .size
378383
384+ def _get_encryption_data_request (self ):
385+ # Save current request cls
386+ download_cls = self ._request_options .pop ('cls' , None )
387+ # Adjust cls for get_properties
388+ self ._request_options ['cls' ] = deserialize_blob_properties
389+
390+ properties = self ._clients .blob .get_properties (** self ._request_options )
391+ # This will return None if there is no encryption metadata or there are parsing errors.
392+ # That is acceptable here, the proper error will be caught and surfaced when attempting
393+ # to decrypt the blob.
394+ self ._encryption_data = parse_encryption_data (properties .metadata )
395+
396+ # Restore cls for download
397+ self ._request_options ['cls' ] = download_cls
398+
379399 def _initial_request (self ):
380400 range_header , range_validation = validate_and_format_range_headers (
381401 self ._initial_range [0 ],
@@ -405,6 +425,9 @@ def _initial_request(self):
405425 # Parse the total file size and adjust the download size if ranges
406426 # were specified
407427 self ._file_size = parse_length_from_content_range (response .properties .content_range )
428+ # Remove any extra encryption data size from blob size
429+ self ._file_size = adjust_blob_size_for_encryption (self ._file_size , self ._encryption_data )
430+
408431 if self ._end_range is not None :
409432 # Use the end range index unless it is over the end of the file
410433 self .size = min (self ._file_size , self ._end_range - self ._start_range + 1 )
@@ -465,7 +488,8 @@ def _initial_request(self):
465488
466489 # If the file is small, the download is complete at this point.
467490 # If file size is large, download the rest of the file in chunks.
468- if response .properties .size != self .size :
491+ # Use less than here for encryption.
492+ if response .properties .size < self .size :
469493 if self ._request_options .get ("modified_access_conditions" ):
470494 self ._request_options ["modified_access_conditions" ].if_match = response .properties .etag
471495 else :
@@ -494,18 +518,25 @@ def chunks(self):
494518 if self ._end_range is not None :
495519 # Use the end range index unless it is over the end of the file
496520 data_end = min (self ._file_size , self ._end_range + 1 )
521+
522+ data_start = self ._initial_range [1 ] + 1 # Start where the first download ended
523+ # For encryption V2 only, adjust start to the end of the fetched data rather than download size
524+ if is_encryption_v2 (self ._encryption_data ):
525+ data_start = (self ._start_range or 0 ) + len (self ._current_content )
526+
497527 iter_downloader = _ChunkDownloader (
498528 client = self ._clients .blob ,
499529 non_empty_ranges = self ._non_empty_ranges ,
500530 total_size = self .size ,
501531 chunk_size = self ._config .max_chunk_get_size ,
502532 current_progress = self ._first_get_size ,
503- start_range = self . _initial_range [ 1 ] + 1 , # start where the first download ended
533+ start_range = data_start ,
504534 end_range = data_end ,
505535 stream = None ,
506536 parallel = False ,
507537 validate_content = self ._validate_content ,
508538 encryption_options = self ._encryption_options ,
539+ encryption_data = self ._encryption_data ,
509540 use_location = self ._location_mode ,
510541 ** self ._request_options
511542 )
@@ -599,18 +630,24 @@ def readinto(self, stream):
599630 # Use the length unless it is over the end of the file
600631 data_end = min (self ._file_size , self ._end_range + 1 )
601632
633+ data_start = self ._initial_range [1 ] + 1 # Start where the first download ended
634+ # For encryption V2 only, adjust start to the end of the fetched data rather than download size
635+ if is_encryption_v2 (self ._encryption_data ):
636+ data_start = (self ._start_range or 0 ) + len (self ._current_content )
637+
602638 downloader = _ChunkDownloader (
603639 client = self ._clients .blob ,
604640 non_empty_ranges = self ._non_empty_ranges ,
605641 total_size = self .size ,
606642 chunk_size = self ._config .max_chunk_get_size ,
607643 current_progress = self ._first_get_size ,
608- start_range = self . _initial_range [ 1 ] + 1 , # Start where the first download ended
644+ start_range = data_start ,
609645 end_range = data_end ,
610646 stream = stream ,
611647 parallel = parallel ,
612648 validate_content = self ._validate_content ,
613649 encryption_options = self ._encryption_options ,
650+ encryption_data = self ._encryption_data ,
614651 use_location = self ._location_mode ,
615652 progress_hook = self ._progress_hook ,
616653 ** self ._request_options
0 commit comments