StarPilot/system/athena/registration.py

#!/usr/bin/env python3
import time
from pathlib import Path

from openpilot.common.params import Params
from openpilot.system.hardware import HARDWARE, PC
from openpilot.system.hardware.hw import Paths
from openpilot.common.swaglog import cloudlog


UNREGISTERED_DONGLE_ID = "UnregisteredDevice"
OFFROAD_UNREGISTERED_ALERT = "Offroad_UnregisteredHardware"
KEYS = {"id_rsa": "RS256",
        "id_ecdsa": "ES256"}


def api_get(*args, **kwargs):
  from openpilot.common.api import api_get as api_get_impl

  return api_get_impl(*args, **kwargs)


def get_key_pair() -> tuple[str, str, str] | tuple[None, None, None]:
  for key, algorithm in KEYS.items():
    private_path = Path(Paths.persist_root()) / "comma" / key
    public_path = Path(Paths.persist_root()) / "comma" / f"{key}.pub"
    if private_path.is_file() and public_path.is_file():
      return algorithm, private_path.read_text(), public_path.read_text()
  return None, None, None


def _has_key_pair() -> bool:
  for key in KEYS:
    private_path = Path(Paths.persist_root()) / "comma" / key
    public_path = Path(Paths.persist_root()) / "comma" / f"{key}.pub"
    if private_path.is_file() and public_path.is_file():
      return True
  return False

def is_registered_device() -> bool:
  dongle = Params().get("DongleId")
  return dongle not in (None, UNREGISTERED_DONGLE_ID)


def _set_registration_alert(params: Params, dongle_id: str | None) -> None:
  show_alert = (dongle_id == UNREGISTERED_DONGLE_ID) and not PC
  if show_alert:
    from openpilot.selfdrive.selfdrived.alertmanager import set_offroad_alert

    set_offroad_alert(OFFROAD_UNREGISTERED_ALERT, True)
  else:
    params.remove(OFFROAD_UNREGISTERED_ALERT)


def register(show_spinner=False, register_konik=False) -> str | None:
  """
  All devices built since March 2024 come with all
  info stored in /persist/. This is kept around
  only for devices built before then.

  With a backend update to take serial number instead
  of dongle ID to some endpoints, this can be removed
  entirely.
  """
  params = Params()

  dongle_id: str | None = params.get("DongleId")
  if dongle_id is None and Path(Paths.persist_root()+"/comma/dongle_id").is_file():
    # not all devices will have this; added early in comma 3X production (2/28/24)
    with open(Paths.persist_root()+"/comma/dongle_id") as f:
      dongle_id = f.read().strip()

  if not _has_key_pair() and not register_konik:
    dongle_id = UNREGISTERED_DONGLE_ID
    cloudlog.warning("missing public key")
  elif dongle_id is None or register_konik:
    import json
    import jwt
    from datetime import datetime, timedelta, UTC
    from openpilot.common.spinner import Spinner

    # Create registration token, in the future, this key will make JWTs directly
    jwt_algo, private_key, public_key = get_key_pair()

    if show_spinner:
      spinner = Spinner()
      spinner.update("registering device")

    # Block until we get the imei
    serial = HARDWARE.get_serial()
    start_time = time.monotonic()
    imei1: str | None = None
    imei2: str | None = None
    while imei1 is None and imei2 is None:
      try:
        imei1, imei2 = HARDWARE.get_imei(0), HARDWARE.get_imei(1)
      except Exception:
        cloudlog.exception("Error getting imei, trying again...")
        time.sleep(1)

      if time.monotonic() - start_time > 60 and show_spinner:
        spinner.update(f"registering device - serial: {serial}, IMEI: ({imei1}, {imei2})")

    backoff = 0
    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)
        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)

        if resp.status_code in (402, 403):
          cloudlog.info(f"Unable to register device, got {resp.status_code}")
          dongle_id = UNREGISTERED_DONGLE_ID
        else:
          dongleauth = json.loads(resp.text)
          dongle_id = dongleauth["dongle_id"]
        break
      except Exception:
        cloudlog.exception("failed to authenticate")
        backoff = min(backoff + 1, 15)
        time.sleep(backoff)

      if time.monotonic() - start_time > 60 and show_spinner:
        spinner.update(f"registering device - serial: {serial}, IMEI: ({imei1}, {imei2})")

    if show_spinner:
      spinner.close()

  if not register_konik and dongle_id != params.get("KonikDongleId"):
    params.put("DongleId", dongle_id)
    params.put("StockDongleId", dongle_id)
    _set_registration_alert(params, dongle_id)
  return dongle_id


if __name__ == "__main__":
  print(register())