1+ import ssl
12from dataclasses import dataclass
23from pathlib import Path
34from typing import Dict , Optional
5+ from urllib .parse import urlparse
46
7+ import aiohttp
8+ import click
59from jumpstarter_driver_composite .client import CompositeClient
610from jumpstarter_driver_opendal .client import FlasherClient , operator_for_path
711from jumpstarter_driver_power .client import PowerClient
@@ -22,10 +26,48 @@ def __post_init__(self):
2226 def boot_to_fastboot (self ):
2327 return self .call ("boot_to_fastboot" )
2428
25- def _upload_file_if_needed (self , file_path : str , operator : Operator | None = None ) -> str :
29+ def _is_http_url (self , path : str ) -> bool :
30+ """Check if the path is an HTTP or HTTPS URL."""
31+ return isinstance (path , str ) and path .startswith (("http://" , "https://" ))
32+
33+ def _download_http (self , url : str , insecure_tls : bool = False ) -> bytes :
34+ """Download a file from HTTP/HTTPS URL using aiohttp."""
35+
36+ async def _download ():
37+ # For http:// or when insecure_tls is set, disable SSL verification
38+ parsed = urlparse (url )
39+ if parsed .scheme == "http" or insecure_tls :
40+ ssl_context : ssl .SSLContext | bool = False
41+ else :
42+ ssl_context = True
43+
44+ connector = aiohttp .TCPConnector (ssl = ssl_context )
45+ async with aiohttp .ClientSession (connector = connector ) as session :
46+ async with session .get (url ) as response :
47+ response .raise_for_status ()
48+ return await response .read ()
49+
50+ return self .portal .call (_download )
51+
52+ def _upload_file_if_needed (
53+ self , file_path : str , operator : Operator | None = None , insecure_tls : bool = False
54+ ) -> str :
2655 if not file_path or not file_path .strip ():
2756 raise ValueError ("File path cannot be empty. Please provide a valid file path." )
2857
58+ if self ._is_http_url (file_path ) and operator is None :
59+ parsed = urlparse (file_path )
60+ is_insecure_http = parsed .scheme == "http"
61+
62+ # usse aiohttp for: http:// URLs, or https:// with insecure_tls
63+ if is_insecure_http or insecure_tls :
64+ filename = Path (parsed .path ).name
65+ self .logger .info (f"Downloading { file_path } to storage as { filename } " )
66+ content = self ._download_http (file_path , insecure_tls = insecure_tls )
67+ self .storage .write_bytes (filename , content )
68+ return filename
69+
70+ # use opendal for local files, https:// (secure), and other schemes
2971 if operator is None :
3072 path_buf , operator , operator_scheme = operator_for_path (file_path )
3173 else :
@@ -46,12 +88,18 @@ def _upload_file_if_needed(self, file_path: str, operator: Operator | None = Non
4688
4789 return filename
4890
49- def flash_images (self , partitions : Dict [str , str ], operators : Optional [Dict [str , Operator ]] = None ):
91+ def flash_images (
92+ self ,
93+ partitions : Dict [str , str ],
94+ operators : Optional [Dict [str , Operator ]] = None ,
95+ insecure_tls : bool = False ,
96+ ):
5097 """Flash images to specified partitions
5198
5299 Args:
53100 partitions: Dictionary mapping partition names to file paths
54101 operators: Optional dictionary mapping partition names to operators
102+ insecure_tls: Skip TLS certificate verification for HTTPS URLs
55103 """
56104 if not partitions :
57105 raise ValueError ("At least one partition must be provided" )
@@ -62,7 +110,7 @@ def flash_images(self, partitions: Dict[str, str], operators: Optional[Dict[str,
62110 for partition , file_path in partitions .items ():
63111 self .logger .info (f"Processing { partition } image: { file_path } " )
64112 operator = operators .get (partition )
65- remote_files [partition ] = self ._upload_file_if_needed (file_path , operator )
113+ remote_files [partition ] = self ._upload_file_if_needed (file_path , operator , insecure_tls = insecure_tls )
66114
67115 self .logger .info ("Checking for fastboot devices on Exporter..." )
68116 detection_result = self .call ("detect_fastboot_device" , 5 , 2.0 )
@@ -84,6 +132,7 @@ def flash(
84132 target : str | None = None ,
85133 operator : Operator | Dict [str , Operator ] | None = None ,
86134 compression = None ,
135+ insecure_tls : bool = False ,
87136 ):
88137 if isinstance (path , dict ):
89138 partitions = path
@@ -109,7 +158,7 @@ def flash(
109158
110159 self .boot_to_fastboot ()
111160
112- result = self .flash_images (partitions , operators )
161+ result = self .flash_images (partitions , operators , insecure_tls = insecure_tls )
113162
114163 self .logger .info ("flash operation completed successfully" )
115164
@@ -130,7 +179,35 @@ def base():
130179 pass
131180
132181 for name , cmd in generic_cli .commands .items ():
133- base .add_command (cmd , name = name )
182+ if name != "flash" :
183+ base .add_command (cmd , name = name )
184+
185+ @base .command ()
186+ @click .argument ("file" , nargs = - 1 , required = False )
187+ @click .option (
188+ "--target" ,
189+ "-t" ,
190+ "target_specs" ,
191+ multiple = True ,
192+ help = "name:file" ,
193+ )
194+ @click .option ("--insecure-tls" , is_flag = True , help = "Skip TLS certificate verification" )
195+ def flash (file , target_specs , insecure_tls ):
196+ """Flash image to DUT"""
197+ if target_specs :
198+ mapping : dict [str , str ] = {}
199+ for spec in target_specs :
200+ if ":" not in spec :
201+ raise click .ClickException (f"Invalid target spec '{ spec } ', expected name:file" )
202+ name , img = spec .split (":" , 1 )
203+ mapping [name ] = img
204+ self .flash (mapping , insecure_tls = insecure_tls )
205+ return
206+
207+ if not file :
208+ raise click .ClickException ("FILE argument is required unless --target/-t is used" )
209+
210+ self .flash (file [0 ], target = None , insecure_tls = insecure_tls )
134211
135212 @base .command ()
136213 def boot_to_fastboot ():
0 commit comments