Merge pull request #26 from onekey-sec/firmware-upload-analysis-configuration-support

firmware upload: add analysis configuration selection

Author: martonilles

README.md

@@ -107,6 +107,17 @@ query {
 """
 res = client.query(GET_PRODUCT_GROUPS)
 default_product_group = next(pg for pg in res["allProductGroups"] if pg["name"] == "Default")
+
+GET_ANALYSIS_CONFIGURATIONS = """
+query {
+  allAnalysisConfigurations {
+    id
+    name
+  }
+}
+"""
+res = client.query(GET_ANALYSIS_CONFIGURATIONS)
+default_analysis_configuration = next(conf for conf in res["allAnalysisConfigurations"] if conf["name"] == "Default")
 ```
 
 You can upload firmwares:
@@ -117,6 +128,7 @@ metadata = FirmwareMetadata(
     vendor_name="myVendor",
     product_name="myProduct",
     product_group_id=default_product_group["id"],
+    analysis_configuration_id=default_analysis_configuration["id"],
 )
 
 firmware_path = Path("/path/to/firmware.bin")

examples/upload_firmware.py

@@ -32,12 +32,24 @@
 response = client.query(GET_ALL_PRODUCT_GROUP_IDS)
 product_group_ids = [pg["id"] for pg in response["allProductGroups"]]
 
+GET_ANALYSIS_CONFIGURATIONS = """
+query {
+  allAnalysisConfigurations {
+    id
+    name
+  }
+}
+"""
+response = client.query(GET_ANALYSIS_CONFIGURATIONS)
+analysis_configuration_ids = [c["id"] for c in response["allAnalysisConfigurations"]]
+
 
 metadata = FirmwareMetadata(
     name="myFirmware",
     vendor_name="myVendor",
     product_name="myProduct",
     product_group_id=product_group_ids[0],
+    analysis_configuration_id=analysis_configuration_ids[0],
 )
 
 firmware_path = Path(sys.argv[2])

onekey_client/cli/firmware_upload.py

@@ -23,6 +23,14 @@
     required=True,
     help="Product group name to add the firmware",
 )
+@click.option(
+    "--analysis-configuration",
+    "analysis_configuration_name",
+    default="Default",
+    show_default=True,
+    required=True,
+    help="Analysis configuration name",
+)
 @click.option("--version", help="Firmware version")
 @click.option("--name", help="Firmware name")
 @click.argument("filename", type=click.Path(exists=True, path_type=Path))
@@ -32,22 +40,17 @@ def upload_firmware(
     product_name: str,
     vendor_name: str,
     product_group_name: str,
+    analysis_configuration_name: str,
     version: Optional[str],
     name: Optional[str],
     filename: Path,
 ):
     """Uploads a firmware to the ONEKEY platform"""
 
-    product_groups = client.get_product_groups()
-
-    try:
-        product_group_id = product_groups[product_group_name]
-    except KeyError:
-        click.echo(f"Missing product group: {product_group_name}")
-        click.echo("Available product groups:")
-        for pg in product_groups.keys():
-            click.echo(f"- {pg}")
-        sys.exit(10)
+    product_group_id = _get_product_group_id_by_name(client, product_group_name)
+    analysis_configuration_id = _get_analysis_configuration_id_by_name(
+        client, analysis_configuration_name
+    )
 
     if name is None:
         name = (
@@ -62,6 +65,7 @@ def upload_firmware(
         product_name=product_name,
         product_group_id=product_group_id,
         version=version,
+        analysis_configuration_id=analysis_configuration_id,
     )
 
     try:
@@ -72,3 +76,31 @@ def upload_firmware(
         for error in e._errors:
             click.echo(f"- {error['message']}")
         sys.exit(11)
+
+
+def _get_product_group_id_by_name(client: Client, product_group_name: str):
+    product_groups = client.get_product_groups()
+
+    try:
+        return product_groups[product_group_name]
+    except KeyError:
+        click.echo(f"Missing product group: {product_group_name}")
+        click.echo("Available product groups:")
+        for pg in product_groups.keys():
+            click.echo(f"- {pg}")
+        sys.exit(10)
+
+
+def _get_analysis_configuration_id_by_name(
+    client: Client, analysis_configuration_name: str
+):
+    analysis_configurations = client.get_analysis_configurations()
+
+    try:
+        return analysis_configurations[analysis_configuration_name]
+    except KeyError:
+        click.echo(f"Missing analysis configuration {analysis_configuration_name}")
+        click.echo("Available analysis configurations:")
+        for config in analysis_configurations.keys():
+            click.echo(f"- {config}")
+        sys.exit(12)

onekey_client/client.py

@@ -209,6 +209,7 @@ def upload_firmware(
                 "releaseDate": metadata.release_date,
                 "notes": metadata.notes,
                 "enableMonitoring": enable_monitoring,
+                "analysisConfigurationId": str(metadata.analysis_configuration_id),
             },
             "vendorName": metadata.vendor_name,
             "productName": metadata.product_name,
@@ -234,6 +235,14 @@ def get_product_groups(self):
         response = self.query(product_groups_query)
         return {pg["name"]: pg["id"] for pg in response["allProductGroups"]}
 
+    @_tenant_required
+    def get_analysis_configurations(self):
+        analysis_configurations_query = load_query(
+            "get_analysis_configurations.graphql"
+        )
+        response = self.query(analysis_configurations_query)
+        return {c["name"]: c["id"] for c in response["allAnalysisConfigurations"]}
+
     def logout(self):
         del self._state
         gc.collect()

onekey_client/models.py

@@ -18,3 +18,4 @@ class FirmwareMetadata(BaseModel):
     product_name: str
     product_category: Optional[str] = None
     product_group_id: UUID
+    analysis_configuration_id: UUID

onekey_client/queries/get_analysis_configurations.graphql

@@ -0,0 +1,6 @@
+query {
+    allAnalysisConfigurations {
+        id
+        name
+    }
+}