1+ #!/usr/bin/env python3
2+ import argparse
3+ import yaml
4+ import os
5+ import urllib .request
6+ import tempfile
7+ import hashlib
8+ import shutil
9+ import re
10+ PRECENT_PROGRESS_SIZE = 5
11+
12+ class ChecksumFailException (Exception ):
13+ pass
14+
15+ IMAGES_CONFIG = os .path .join (os .path .dirname (__file__ ), "images.yml" )
16+ RETRY = 3
17+
18+ def ensure_dir (d , chmod = 0o777 ):
19+ """
20+ Ensures a folder exists.
21+ Returns True if the folder already exists
22+ """
23+ if not os .path .exists (d ):
24+ os .makedirs (d , chmod )
25+ os .chmod (d , chmod )
26+ return False
27+ return True
28+
29+ def read_images ():
30+ if not os .path .isfile (IMAGES_CONFIG ):
31+ raise Exception (f"Error: Remotes config file not found: { IMAGES_CONFIG } )
32+ with open (IMAGES_CONFIG ,'r' ) as f :
33+ output = yaml .safe_load (f )
34+ return output
35+
36+ class DownloadProgress :
37+ last_precent : float = 0
38+ def show_progress (self , block_num , block_size , total_size ):
39+ new_precent = round (block_num * block_size / total_size * 100 , 1 )
40+ if new_precent > self .last_precent + PRECENT_PROGRESS_SIZE :
41+ print (f"{ new_precent } , end = "\r " )
42+ self .last_precent = new_precent
43+
44+ def get_file_name (headers ):
45+ return re .findall ("filename=(\S+)" , headers ["Content-Disposition" ])[0 ]
46+
47+ def get_sha256 (filename ):
48+ sha256_hash = hashlib .sha256 ()
49+ with open (filename ,"rb" ) as f :
50+ for byte_block in iter (lambda : f .read (4096 ),b"" ):
51+ sha256_hash .update (byte_block )
52+ file_checksum = sha256_hash .hexdigest ()
53+ return file_checksum
54+ return
55+
56+ def download_image_http (board : str , dest_folder : str , redownload : bool = False ):
57+ url = board ["url" ]
58+ checksum = board ["checksum" ]
59+
60+ with tempfile .TemporaryDirectory () as tmpdirname :
61+ print ('created temporary directory' , tmpdirname )
62+ temp_file_name = os .path .join (tmpdirname , "image.xz" )
63+ temp_file_checksum = os .path .join (tmpdirname , "checksum.sha256" )
64+
65+ for r in range (RETRY ):
66+ try :
67+ # Get sha and confirm its the right image
68+ download_progress = DownloadProgress ()
69+ _ , headers_checksum = urllib .request .urlretrieve (checksum , temp_file_checksum , download_progress .show_progress )
70+ file_name_checksum = get_file_name (headers_checksum )
71+
72+ checksum_data = None
73+ with open (temp_file_checksum , 'r' ) as f :
74+ checksum_data = f .read ()
75+
76+ checksum_data_parsed = [x .strip () for x in checksum_data .split ()]
77+ online_checksum = checksum_data_parsed [0 ]
78+ file_name_from_checksum = checksum_data_parsed [1 ]
79+ dest_file_name = os .path .join (dest_folder , file_name_from_checksum )
80+ print (f"Downloading { dest_file_name } )
81+
82+ if os .path .isfile (dest_file_name ):
83+ file_checksum = get_sha256 (dest_file_name )
84+ if file_checksum == online_checksum :
85+ # We got file and checksum is right
86+ return
87+ # Get the file
88+ download_progress = DownloadProgress ()
89+ _ , headers = urllib .request .urlretrieve (url , temp_file_name , download_progress .show_progress )
90+
91+ file_name = get_file_name (headers )
92+ file_checksum = get_sha256 (temp_file_name )
93+ if file_checksum != online_checksum :
94+ print (f'Failed. Attempt # { r + 1 } { file_checksum } { online_checksum } )
95+ continue
96+ ensure_dir (os .path .dirname (dest_file_name ))
97+ shutil .move (temp_file_name , dest_file_name )
98+
99+ except Exception as e :
100+ if r < 2 :
101+ print (f'Failed. Attempt # { r + 1 } { e } )
102+ else :
103+ print ('Error encoutered at {RETRY} attempt' )
104+ print (e )
105+ else :
106+ print (f"Success: { temp_file_name } )
107+ break
108+ return
109+
110+ if __name__ == "__main__" :
111+ parser = argparse .ArgumentParser (add_help = True , description = 'Download images based on BASE_BOARD and BASE_O' )
112+ parser .add_argument ('WORKSPACE_SUFFIX' , nargs = '?' , default = "default" , help = "The workspace folder suffix used folder" )
113+ parser .add_argument ('-s' , '--sha256' , action = 'store_true' , help = 'Create a sha256 hash for the .img file in .sha256' )
114+ args = parser .parse_args ()
115+
116+ images = read_images ()
117+
118+ base_board = os .environ .get ("BASE_BOARD" , None )
119+ base_image_path = os .environ .get ("BASE_IMAGE_PATH" , None )
120+
121+ if base_board is not None and base_board in images ["images" ]:
122+ if images ["images" ][base_board ]["type" ] == "http" :
123+ download_image_http (images ["images" ][base_board ], base_image_path )
124+ elif images ["images" ][base_board ]["type" ] == "torrent" :
125+ print ("Error: Torrent not implemented" )
126+ exit (1 )
127+ else :
128+ print ("Error: Unsupported image download type" )
129+ exit (1 )
130+
131+ print ("Done" )
0 commit comments