Repository: microsoft/onefuzz
File: src/api-service/__app__/onefuzzlib/azure/image.py (Python)

Code before:
#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from typing import Union

from azure.core.exceptions import ResourceNotFoundError
from memoization import cached
from msrestazure.azure_exceptions import CloudError
from msrestazure.tools import parse_resource_id
from onefuzztypes.enums import OS, ErrorCode
from onefuzztypes.models import Error
from onefuzztypes.primitives import Region

from .compute import get_compute_client


@cached(ttl=60)
def get_os(region: Region, image: str) -> Union[Error, OS]:
    client = get_compute_client()
    parsed = parse_resource_id(image)
    if "resource_group" in parsed:
        if parsed["type"] == "galleries":
            try:
                name = client.gallery_images.get(
                    parsed["resource_group"], parsed["name"], parsed["child_name_1"]
                ).os_type.lower()
            except (ResourceNotFoundError, CloudError) as err:
                return Error(code=ErrorCode.INVALID_IMAGE, errors=[str(err)])
        else:
            try:
                name = client.images.get(
                    parsed["resource_group"], parsed["name"]
                ).storage_profile.os_disk.os_type.lower()
            except (ResourceNotFoundError, CloudError) as err:
                return Error(code=ErrorCode.INVALID_IMAGE, errors=[str(err)])
    else:
        publisher, offer, sku, version = image.split(":")
        try:
            if version == "latest":
                version = client.virtual_machine_images.list(
                    region, publisher, offer, sku, top=1
                )[0].name
            name = client.virtual_machine_images.get(
                region, publisher, offer, sku, version
            ).os_disk_image.operating_system.lower()
        except (ResourceNotFoundError, CloudError) as err:
            return Error(code=ErrorCode.INVALID_IMAGE, errors=[str(err)])
    return OS[name]


Review comment (suggestion):
Thinking through this: I think we've been assuming that if the `"resource_group"` key is present, we'll at least have a _key_ in `parsed` for every return key described in the `parse_resource_id()` docs. Unfortunately, that's not valid, and per the very weak contract of our parameter type, we can actually get an _empty_ dict.

Example:
```python
>>> rid_provider = '/subscriptions/11111111-1111-1111-1111-111111111111/resourceGroups/my-resource-group/providers/Microsoft.Web/sites/my-site-name'
>>> rid_no_provider = '/subscriptions/11111111-1111-1111-1111-111111111111/resourceGroups/my-resource-group'
>>> parse_resource_id(rid_provider)
{'children': '',
 'name': 'my-site-name',
 'namespace': 'Microsoft.Web',
 'resource_group': 'my-resource-group',
 'resource_name': 'my-site-name',
 'resource_namespace': 'Microsoft.Web',
 'resource_parent': '',
 'resource_type': 'sites',
 'subscription': '11111111-1111-1111-1111-111111111111',
 'type': 'sites'}
>>> parse_resource_id(rid_no_provider)
{'resource_group': 'my-resource-group',
 'subscription': '11111111-1111-1111-1111-111111111111'}
```

This is not a great situation, and we should probably add types (in a `Union`) to explicitly represent the different expected cases, and allow clearer dispatching at this site.

Code after:
#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from typing import Union

from azure.core.exceptions import ResourceNotFoundError
from memoization import cached
from msrestazure.azure_exceptions import CloudError
from msrestazure.tools import parse_resource_id
from onefuzztypes.enums import OS, ErrorCode
from onefuzztypes.models import Error
from onefuzztypes.primitives import Region

from .compute import get_compute_client


@cached(ttl=60)
def get_os(region: Region, image: str) -> Union[Error, OS]:
    client = get_compute_client()
    # The dict returned here may not have any defined keys.
    #
    # See: https://github.com/Azure/msrestazure-for-python/blob/v0.6.3/msrestazure/tools.py#L134  # noqa: E501
    parsed = parse_resource_id(image)
    if "resource_group" in parsed:
        if parsed["type"] == "galleries":
            try:
                # See: https://docs.microsoft.com/en-us/rest/api/compute/gallery-images/get#galleryimage  # noqa: E501
                name = client.gallery_images.get(
                    parsed["resource_group"], parsed["name"], parsed["child_name_1"]
                ).os_type.lower()
            except (ResourceNotFoundError, CloudError) as err:
                return Error(code=ErrorCode.INVALID_IMAGE, errors=[str(err)])
        else:
            try:
                # See: https://docs.microsoft.com/en-us/rest/api/compute/images/get
                name = client.images.get(
                    parsed["resource_group"], parsed["name"]
                ).storage_profile.os_disk.os_type.lower()
            except (ResourceNotFoundError, CloudError) as err:
                return Error(code=ErrorCode.INVALID_IMAGE, errors=[str(err)])
    else:
        publisher, offer, sku, version = image.split(":")
        try:
            if version == "latest":
                version = client.virtual_machine_images.list(
                    region, publisher, offer, sku, top=1
                )[0].name
            name = client.virtual_machine_images.get(
                region, publisher, offer, sku, version
            ).os_disk_image.operating_system.lower()