athenad: get ES256 key (#36845)

* fix

* why not format

* fix typing

* cast

common/api.py

@@ -7,9 +7,10 @@ from openpilot.system.version import get_version
 
 API_HOST = os.getenv('API_HOST', 'https://api.commadotai.com')
 
-       # name : jwt signature algorithm
-KEYS = {"id_rsa" : "RS256",
-        "id_ecdsa" : "ES256"}
+# name: jwt signature algorithm
+KEYS = {"id_rsa": "RS256",
+        "id_ecdsa": "ES256"}
+
 
 class Api:
   def __init__(self, dongle_id):
@@ -50,7 +51,8 @@ def api_get(endpoint, method='GET', timeout=None, access_token=None, **params):
 
   return requests.request(method, API_HOST + "/" + endpoint, timeout=timeout, headers=headers, params=params)
 
-def get_key_pair():
+
+def get_key_pair() -> tuple[str, str, str] | tuple[None, None, None]:
   for key in KEYS:
     if os.path.isfile(Paths.persist_root() + f'/comma/{key}') and os.path.isfile(Paths.persist_root() + f'/comma/{key}.pub'):
       with open(Paths.persist_root() + f'/comma/{key}') as private, open(Paths.persist_root() + f'/comma/{key}.pub') as public:

system/athena/athenad.py

@@ -30,7 +30,7 @@ from websocket import (ABNF, WebSocket, WebSocketException, WebSocketTimeoutExce
 import cereal.messaging as messaging
 from cereal import log
 from cereal.services import SERVICE_LIST
-from openpilot.common.api import Api
+from openpilot.common.api import Api, get_key_pair
 from openpilot.common.utils import CallbackReader, get_upload_stream
 from openpilot.common.params import Params
 from openpilot.common.realtime import set_core_affinity
@@ -523,11 +523,8 @@ def startLocalProxy(global_end_event: threading.Event, remote_ws_uri: str, local
 
 @dispatcher.add_method
 def getPublicKey() -> str | None:
-  if not os.path.isfile(Paths.persist_root() + '/comma/id_rsa.pub'):
-    return None
-
-  with open(Paths.persist_root() + '/comma/id_rsa.pub') as f:
-    return f.read()
+  _, _, public_key = get_key_pair()
+  return public_key
 
 
 @dispatcher.add_method

system/athena/registration.py

@@ -2,6 +2,7 @@
 import time
 import json
 import jwt
+from typing import cast
 from pathlib import Path
 
 from datetime import datetime, timedelta, UTC
@@ -69,7 +70,8 @@ def register(show_spinner=False) -> str | None:
     start_time = time.monotonic()
     while True:
       try:
-        register_token = jwt.encode({'register': True, 'exp': datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1)}, private_key, algorithm=jwt_algo)
+        register_token = jwt.encode({'register': True, 'exp': datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1)},
+                                    cast(str, private_key), algorithm=jwt_algo)
         cloudlog.info("getting pilotauth")
         resp = api_get("v2/pilotauth/", method='POST', timeout=15,
                        imei=imei1, imei2=imei2, serial=serial, public_key=public_key, register_token=register_token)