7
7
import hashlib
8
8
import shutil
9
9
import re
10
+ from typing import Dict , Any , Optional , cast
11
+ from common import get_image_config , read_images
10
12
PRECENT_PROGRESS_SIZE = 5
11
13
12
14
class ChecksumFailException (Exception ):
13
15
pass
14
16
15
- IMAGES_CONFIG = os .path .join (os .path .dirname (__file__ ), "images.yml" )
16
17
RETRY = 3
17
18
18
19
def ensure_dir (d , chmod = 0o777 ):
@@ -26,13 +27,6 @@ def ensure_dir(d, chmod=0o777):
26
27
return False
27
28
return True
28
29
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
30
class DownloadProgress :
37
31
last_precent : float = 0
38
32
def show_progress (self , block_num , block_size , total_size ):
@@ -41,8 +35,10 @@ def show_progress(self, block_num, block_size, total_size):
41
35
print (f"{ new_precent } %" , end = "\r " )
42
36
self .last_precent = new_precent
43
37
44
- def get_file_name (headers ):
45
- return re .findall ("filename=(\S+)" , headers ["Content-Disposition" ])[0 ]
38
+ def get_file_name (headers , url ):
39
+ if "Content-Disposition" in headers .keys ():
40
+ return re .findall ("filename=(\S+)" , headers ["Content-Disposition" ])[0 ]
41
+ return url .split ('/' )[- 1 ]
46
42
47
43
def get_sha256 (filename ):
48
44
sha256_hash = hashlib .sha256 ()
@@ -53,7 +49,7 @@ def get_sha256(filename):
53
49
return file_checksum
54
50
return
55
51
56
- def download_image_http (board : str , dest_folder : str , redownload : bool = False ):
52
+ def download_image_http (board : Dict [ str , Any ] , dest_folder : str , redownload : bool = False ):
57
53
url = board ["url" ]
58
54
checksum = board ["checksum" ]
59
55
@@ -67,7 +63,7 @@ def download_image_http(board: str, dest_folder: str, redownload: bool = False):
67
63
# Get sha and confirm its the right image
68
64
download_progress = DownloadProgress ()
69
65
_ , headers_checksum = urllib .request .urlretrieve (checksum , temp_file_checksum , download_progress .show_progress )
70
- file_name_checksum = get_file_name (headers_checksum )
66
+ file_name_checksum = get_file_name (headers_checksum , checksum )
71
67
72
68
checksum_data = None
73
69
with open (temp_file_checksum , 'r' ) as f :
@@ -88,7 +84,7 @@ def download_image_http(board: str, dest_folder: str, redownload: bool = False):
88
84
download_progress = DownloadProgress ()
89
85
_ , headers = urllib .request .urlretrieve (url , temp_file_name , download_progress .show_progress )
90
86
91
- file_name = get_file_name (headers )
87
+ file_name = get_file_name (headers , url )
92
88
file_checksum = get_sha256 (temp_file_name )
93
89
if file_checksum != online_checksum :
94
90
print (f'Failed. Attempt # { r + 1 } , checksum missmatch: { file_checksum } expected: { online_checksum } ' )
@@ -102,6 +98,7 @@ def download_image_http(board: str, dest_folder: str, redownload: bool = False):
102
98
else :
103
99
print ('Error encoutered at {RETRY} attempt' )
104
100
print (e )
101
+ exit (1 )
105
102
else :
106
103
print (f"Success: { temp_file_name } " )
107
104
break
@@ -118,14 +115,25 @@ def download_image_http(board: str, dest_folder: str, redownload: bool = False):
118
115
base_board = os .environ .get ("BASE_BOARD" , None )
119
116
base_image_path = os .environ .get ("BASE_IMAGE_PATH" , None )
120
117
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" :
118
+ if base_image_path is None :
119
+ print (f'Error: did not find image config file' )
120
+ exit (1 )
121
+ cast (str , base_image_path )
122
+
123
+ image_config = get_image_config ()
124
+ if image_config is not None :
125
+ if image_config ["type" ] == "http" :
126
+ print (f"Downloading image for { base_board } " )
127
+ download_image_http (image_config , base_image_path )
128
+ elif image_config ["type" ] == "torrent" :
125
129
print ("Error: Torrent not implemented" )
126
130
exit (1 )
127
131
else :
128
- print (" Error: Unsupported image download type" )
132
+ print (f' Error: Unsupported image download type: { image_config [ "type" ] } ' )
129
133
exit (1 )
134
+ else :
135
+ print (f"Error: Image config not found for: { base_board } " )
136
+ exit (1 )
137
+
130
138
131
- print ("Done" )
139
+ print ("Done" )
0 commit comments