11#!/usr/bin/env python3
2- # Copyright (c) 2023 , Arm Limited or its affiliates. All rights reserved.
2+ # Copyright (c) 2024 , Arm Limited or its affiliates. All rights reserved.
33# SPDX-License-Identifier : Apache-2.0
44#
55# Licensed under the Apache License, Version 2.0 (the "License");
1919
2020import subprocess
2121import re
22+ import hashlib
23+ import threading
24+ import os
2225
2326# Note: Precious partitions dictionary, this is a set of partition types which might have firmware
2427# and to be refrained from read/write operations. The list is not exhaustive might see additions in future.
2528precious_parts_mbr = {
2629 "Protective partition" :"0xF8" ,
2730 "EFI system partition" :"0xEF"
28-
2931}
3032
3133precious_parts_gpt = {
3436 "U-Boot environment partition" :"3DE21764-95BD-54BD-A5C3-4ABE786F38A8"
3537}
3638
39+ def input_with_timeout (prompt , timeout = 5 ):
40+ print (prompt , end = '' , flush = True )
41+ input_queue = []
42+ def get_user_input ():
43+ try :
44+ input_queue .append (input ())
45+ except EOFError :
46+ pass
47+ user_input_thread = threading .Thread (target = get_user_input )
48+ user_input_thread .daemon = True
49+ user_input_thread .start ()
50+ user_input_thread .join (timeout )
51+ return input_queue [0 ] if input_queue else "no"
52+
53+ def calculate_sha256 (file_path ):
54+ sha256_hash = hashlib .sha256 ()
55+ with open (file_path , "rb" ) as f :
56+ # Read and update hash string value in blocks of 4K
57+ for byte_block in iter (lambda : f .read (4096 ), b"" ):
58+ sha256_hash .update (byte_block )
59+ return sha256_hash .hexdigest ()
60+
61+
62+ def get_partition_space (partition_path ): #to calculate used blocks and available blocks
63+ command = f"df -B 512 { partition_path }
64+ result = subprocess .run (command , shell = True , text = True , check = True , capture_output = True )
65+ lines = result .stdout .strip ().split ('\n ' )
66+ if len (lines ) > 1 :
67+ parts = lines [1 ].split ()
68+ used_blocks = int (parts [0 ])
69+ available_blocks = int (parts [1 ])
70+ return used_blocks , available_blocks
71+ else :
72+ print (f"WARNING: Unable to parse partition space for { partition_path } )
73+ return 0 , 0 # Default to 0 if parsing fails
74+
75+ def perform_write_check (partition_label , partition_id , precious_parts ):
76+
77+ # Check if partition is precious
78+ user_input = 'no' if partition_id in precious_parts .values () else input_with_timeout (
79+ f"Do you want to perform a write check on /dev/{ partition_label } , 5 ).lower ()
80+
81+ if user_input == 'yes' and partition_id not in precious_parts .values ():
82+ used_blocks , available_blocks = get_partition_space (f"/dev/{ partition_label } )
83+ if available_blocks > 0 :
84+
85+ # Prepare hello.txt content with padding to 512 bytes
86+ hello_content = "Hello!" .ljust (512 , '\x00 ' ) # Pad the content to 512 bytes
87+ with open ("hello.txt" , "wb" ) as f :
88+ f .write (hello_content .encode ('utf-8' ))
89+ original_sha256 = calculate_sha256 ("hello.txt" )
90+
91+ # Backup filename specific to partition
92+ backup_filename = f"{ partition_label }
93+ # backup command
94+ print ("INFO: Creating backup of the current block before write check..." )
95+ backup_command = f"dd if=/dev/{ partition_label } { backup_filename } { used_blocks }
96+ subprocess .run (backup_command , shell = True , check = True )
97+
98+ # Write padded hello.txt to the device
99+ print ("INFO: Writing test data to the device for write check..." )
100+ write_command = f"dd if=hello.txt of=/dev/{ partition_label } { used_blocks }
101+ subprocess .run (write_command , shell = True , check = True )
102+
103+ # Read back the 512-byte block
104+ read_back_file = "read_hello.txt"
105+ print ("INFO: Reading back the test data for verification..." )
106+ read_command = f"dd if=/dev/{ partition_label } { read_back_file } { used_blocks }
107+ subprocess .run (read_command , shell = True , check = True )
108+
109+ # Calculate SHA256 for the padded content to ensure a fair comparison
110+ read_back_sha256 = calculate_sha256 (read_back_file )
111+
112+ # Verify checksums
113+ print (f"Original SHA256: { original_sha256 } )
114+ print (f"Read-back SHA256: { read_back_sha256 } )
115+ if original_sha256 == read_back_sha256 :
116+ print (f"INFO: write check passed on /dev/{ partition_label } )
117+
118+ # Restore the backup
119+ print ("INFO: Restoring the backup to the device after write check..." )
120+ restore_command = f"dd if={ backup_filename } { partition_label } { used_blocks }
121+ subprocess .run (restore_command , shell = True , check = True )
122+ print (f"INFO: Backup restored for /dev/{ partition_label } )
123+ else :
124+ print (f"WARNING: Data integrity check failed for /dev/{ partition_label } )
125+
126+ # Clean up
127+ os .remove ("hello.txt" )
128+ os .remove (read_back_file )
129+ os .remove (backup_filename )
130+
131+ else :
132+ print (f"WARNING: No available space for write check on /dev/{ partition_label } )
133+
134+
37135if __name__ == "__main__" :
38136 try :
39137 # find all disk block devices
55153 print (f"INFO: Block device : /dev/{ disk } )
56154
57155 # check whether disk uses MBR or GPT partition table
58- command = f"timeout 5 gdisk -l /dev/{ disk }
156+ command = f"timeout 10 gdisk -l /dev/{ disk }
59157 result = subprocess .run (command , shell = True , text = True , check = False , capture_output = True )
60158
61159 if "MBR: MBR only" in result .stdout :
76174
77175 # skip if no partitions
78176 if num_parts == 0 :
79- print (f"INFO: No partitions detected for { disk } )
177+ print (f"INFO: No partitions detected for { disk } /write ..." )
80178 continue
81179
82180 # get partition labels
83181 command = f"lsblk /dev/{ disk }
84182 result = subprocess .run (command , shell = True , text = True , check = False , capture_output = True )
85183 part_lables = re .findall (r'[├└]─(\S+)' , result .stdout )
86184
87-
88185 # if disk follows MBR
89186 if part_table == "MBR" :
90187 table_header_row = ["Device" , "Boot Start" , "End Sectors" , "Size" , "Id" , "Type" ]
91188 command = f"fdisk -l /dev/{ disk }
92189 result = subprocess .run (command , shell = True , text = True , check = False , capture_output = True )
93190
94191 # get MBR partition type Ids
95- lines = ( result .stdout ) .strip ().split ('\n ' )
192+ lines = result .stdout .strip ().split ('\n ' )
96193 collect_lines = False
97194 mbr_part_ids = []
98195
99- # Iterate through the lines and collect lines after the header line
100196 for line in lines :
101197 if collect_lines :
102- # Split the line into columns using spaces and extract the 6th column
103198 columns = line .split ()
104199 if len (columns ) >= 6 :
105- mbr_part_ids .append ("0x" + ( columns [5 ]) .upper ())
200+ mbr_part_ids .append ("0x" + columns [5 ].upper ())
106201 elif all (substring in line for substring in table_header_row ):
107202 collect_lines = True
108203
109- # check if parsing went well for MBR
110204 if not (len (part_lables ) == len (mbr_part_ids ) == num_parts ):
111205 print (f"INFO: Error parsing MBR partition Ids/partition labels/number of partition(s) for { disk } )
112206
113- # iterate partitions
114207 for index in range (0 , num_parts ):
115208 print (f"\n INFO: Partition : /dev/{ part_lables [index ]} { mbr_part_ids [index ]} )
116- # check if the partition is precious
209+
117210 if mbr_part_ids [index ] in precious_parts_mbr .values ():
118211 for key , value in precious_parts_mbr .items ():
119212 if value == mbr_part_ids [index ]:
120213 print (f"INFO: { part_lables [index ]} )
214+ used_blocks = get_partition_space (f"/dev/{ part_lables [index ]} )
215+ print (f"INFO: Number of 512B blocks used on /dev/{ part_lables [index ]} { used_blocks } )
121216 print (f" { key } { value } )
122- print (" Skipping block read..." )
217+ print (" Skipping block read/write ..." )
123218 continue
124- # perform block read of 1 block (1MB)
125219 else :
126- command = f"dd if=/dev/{ part_lables [index ]}
220+ command = f"dd if=/dev/{ part_lables [index ]} > /dev/null "
127221 print (f"INFO: Performing block read on /dev/{ part_lables [index ]} { mbr_part_ids [index ]} )
128222 result = subprocess .run (command , shell = True , text = True , check = False , capture_output = True )
129-
130223 if result .returncode == 0 :
131224 print (f"INFO: Block read on /dev/{ part_lables [index ]} { mbr_part_ids [index ]} )
225+ # Call the perform_write_check function for the partition
226+ perform_write_check (part_lables [index ], partition_guid_code , precious_parts_gpt )
132227 else :
133228 print (f"INFO: Block read on /dev/{ part_lables [index ]} { mbr_part_ids [index ]} )
134229 print ("\n ********************************************************************************************************************************\n " )
172267
173268 # skip block read if "Platform required" bit is set
174269 if lsb == 1 :
175- print (f"INFO: Platform required attribute set for { part_lables [index ]} )
270+ print (f"INFO: Platform required attribute set for { part_lables [index ]} /write ..." )
176271 continue
177272 # check if the partition is precious
178273 if partition_guid_code in precious_parts_gpt .values ():
274+
275+ used_blocks = get_partition_space (f"/dev/{ part_lables [index ]} )
276+
179277 for key , value in precious_parts_gpt .items ():
180278 if value == partition_guid_code :
181279 print (f"INFO: { part_lables [index ]} )
280+ print (f"INFO: Number of 512B blocks used on /dev/{ part_lables [index ]} { used_blocks } )
182281 print (f" { key } { value } )
183- print (" Skipping block read..." )
282+ print (" Skipping block read/write ..." )
184283 continue
185284 # perform block read of 1 block (1MB)
186285 else :
190289
191290 if result .returncode == 0 :
192291 print (f"INFO: Block read on /dev/{ part_lables [index ]} { partition_guid_code } )
292+ # Call the perform_write_check function for the partition
293+ perform_write_check (part_lables [index ], partition_guid_code , precious_parts_gpt )
193294 else :
194295 print (f"INFO: Block read on /dev/{ part_lables [index ]} { partition_guid_code } )
296+
195297 print ("\n ********************************************************************************************************************************\n " )
196298 else :
197299 print (f"INFO: Invalid partition table, expected MBR or GPT reported type = { part_table } )
198300
199301 except Exception as e :
200302 print (f"Error occurred: { e } )
201- exit (1 )
303+ exit (1 )
0 commit comments