55""" 
66
77
8- import  sys 
98from  ..common  import  Chain 
109from  ..errors  import  (
1110    DEVICE_NOT_INITIALIZED ,
1211    DeviceNotReadyError ,
1312    common_err_msgs ,
1413    handle_errors ,
1514)
16- from  .trezorlib  import  protobuf ,  debuglink 
15+ from  .trezorlib  import  protobuf 
1716from  .trezorlib .transport  import  (
1817    udp ,
1918    webusb ,
3534    Optional ,
3635    Sequence ,
3736)
37+ import  copy 
3838
3939py_enumerate  =  enumerate   # Need to use the enumerate built-in but there's another function already named that 
4040
@@ -223,8 +223,8 @@ def __init__(
223223        self .serial_no  =  serial_no 
224224        self .boardloader_version  =  boardloader_version 
225225
226- 
227- DEFAULT_MAPPING .register (OnekeyFeatures )
226+ ONEKEY_MAPPING   =   copy . deepcopy ( DEFAULT_MAPPING ) 
227+ ONEKEY_MAPPING .register (OnekeyFeatures )
228228
229229USB_IDS  =  {(0x1209 , 0x4F4A ), (0x1209 , 0x4F4B ), }
230230
@@ -233,15 +233,15 @@ def __init__(
233233    minimum_version = (2 , 11 , 0 ),
234234    vendors = VENDORS ,
235235    usb_ids = USB_IDS ,
236-     default_mapping = DEFAULT_MAPPING ,
236+     default_mapping = ONEKEY_MAPPING ,
237237)
238238
239239ONEKEY_TOUCH  =  TrezorModel (
240240    name = "T" ,
241241    minimum_version = (4 , 2 , 0 ),
242242    vendors = VENDORS ,
243243    usb_ids = USB_IDS ,
244-     default_mapping = DEFAULT_MAPPING ,
244+     default_mapping = ONEKEY_MAPPING ,
245245)
246246
247247ONEKEYS  =  (ONEKEY_LEGACY , ONEKEY_TOUCH )
@@ -256,29 +256,14 @@ def model_by_name(name: str) -> Optional[TrezorModel]:
256256
257257# ===============overwrite methods for onekey device begin============ 
258258
259+ def  retrieval_version (self : object ):
260+     version  =  (* map (int , self .features .onekey_version .split ("." )), )
261+     return  version 
259262
260- def  _refresh_features (self : object , features : Features ) ->  None :
261-     """Update internal fields based on passed-in Features message.""" 
262-     if  not  self .model :
263-         self .model  =  model_by_name (features .model  or  "1" )
264-         if  self .model  is  None :
265-             raise  RuntimeError ("Unsupported OneKey model" )
266- 
267-     if  features .vendor  not  in self .model .vendors :
268-         raise  RuntimeError ("Unsupported device" )
269-     self .features  =  features 
270-     self .version  =  (* map (int , self .features .onekey_version .split ("." )), )
271-     self .check_firmware_version (warn_only = True )
272-     if  self .features .session_id  is  not None :
273-         self .session_id  =  self .features .session_id 
274-         self .features .session_id  =  None 
275- 
276- def  button_request (self : object , code : Optional [int ]) ->  None :
277-     if  not  self .prompt_shown :
278-         print ("Please confirm action on your OneKey device" , file = sys .stderr )
279-     if  not  self .always_prompt :
280-         self .prompt_shown  =  True 
281- 
263+ def  ensure_model (self : object , features ):
264+     assert  self .model  is  not None , "Unsupported OneKey model" 
265+     # Correct the correct model 
266+     self .model  =  model_by_name (features .model  or  "1" )
282267
283268# ===============overwrite methods for onekey device end============ 
284269
@@ -291,12 +276,9 @@ def __init__(
291276        expert : bool  =  False ,
292277        chain : Chain  =  Chain .MAIN ,
293278    ) ->  None :
294-         super ().__init__ (path , password , expert , chain , webusb_ids = USB_IDS , sim_path = ONEKEY_EMULATOR_PATH )
295-         self .client ._refresh_features  =  MethodType (_refresh_features , self .client )
296-         if  not  isinstance (self .client .ui , debuglink .DebugUI ):
297-             self .client .ui .button_request  =  MethodType (button_request , self .client .ui )
298-         self .type  =  "OneKey" 
299- 
279+         super ().__init__ (path , password , expert , chain , webusb_ids = USB_IDS , sim_path = ONEKEY_EMULATOR_PATH , model = ONEKEY_LEGACY , device_type = "OneKey" )
280+         self .client .retrieval_version  =  MethodType (retrieval_version , self .client )
281+         self .client .ensure_model  =  MethodType (ensure_model , self .client )
300282
301283def  enumerate (
302284    password : Optional [str ] =  None , expert : bool  =  False , chain : Chain  =  Chain .MAIN 
0 commit comments