diff --git a/src/api-service/__app__/agent_registration/__init__.py b/src/api-service/__app__/agent_registration/__init__.py index fe6dc4a825..b7b79f45ff 100644 --- a/src/api-service/__app__/agent_registration/__init__.py +++ b/src/api-service/__app__/agent_registration/__init__.py @@ -12,9 +12,9 @@ from onefuzztypes.requests import AgentRegistrationGet, AgentRegistrationPost from onefuzztypes.responses import AgentRegistration -from ..onefuzzlib.azure.containers import StorageType from ..onefuzzlib.azure.creds import get_instance_url from ..onefuzzlib.azure.queue import get_queue_sas +from ..onefuzzlib.azure.storage import StorageType from ..onefuzzlib.endpoint_authorization import call_if_agent from ..onefuzzlib.pools import Node, NodeMessage, NodeTasks, Pool from ..onefuzzlib.request import not_ok, ok, parse_uri diff --git a/src/api-service/__app__/containers/__init__.py b/src/api-service/__app__/containers/__init__.py index 3c5be88d80..4bd86236c4 100644 --- a/src/api-service/__app__/containers/__init__.py +++ b/src/api-service/__app__/containers/__init__.py @@ -13,13 +13,13 @@ from onefuzztypes.responses import BoolResult, ContainerInfo, ContainerInfoBase from ..onefuzzlib.azure.containers import ( - StorageType, create_container, delete_container, get_container_metadata, get_container_sas_url, get_containers, ) +from ..onefuzzlib.azure.storage import StorageType from ..onefuzzlib.endpoint_authorization import call_if_user from ..onefuzzlib.request import not_ok, ok, parse_request diff --git a/src/api-service/__app__/download/__init__.py b/src/api-service/__app__/download/__init__.py index 5fe9566c04..95390ccb06 100644 --- a/src/api-service/__app__/download/__init__.py +++ b/src/api-service/__app__/download/__init__.py @@ -8,11 +8,11 @@ from onefuzztypes.models import Error, FileEntry from ..onefuzzlib.azure.containers import ( - StorageType, blob_exists, container_exists, get_file_sas_url, ) +from ..onefuzzlib.azure.storage import StorageType from ..onefuzzlib.endpoint_authorization import call_if_user from ..onefuzzlib.request import not_ok, parse_uri, redirect diff --git a/src/api-service/__app__/onefuzzlib/azure/containers.py b/src/api-service/__app__/onefuzzlib/azure/containers.py index 4e0077b701..1b52701139 100644 --- a/src/api-service/__app__/onefuzzlib/azure/containers.py +++ b/src/api-service/__app__/onefuzzlib/azure/containers.py @@ -4,81 +4,134 @@ # Licensed under the MIT License. import datetime +import logging import os import urllib.parse -from enum import Enum -from typing import Any, Dict, Optional, Union, cast +from typing import Dict, Optional, Union, cast from azure.common import AzureHttpError, AzureMissingResourceHttpError -from azure.storage.blob import BlobPermissions, ContainerPermissions +from azure.storage.blob import BlobPermissions, BlockBlobService, ContainerPermissions from memoization import cached +from onefuzztypes.primitives import Container -from .creds import get_blob_service, get_func_storage, get_fuzz_storage +from .storage import ( + StorageType, + choose_account, + get_accounts, + get_storage_account_name_key, +) -class StorageType(Enum): - corpus = "corpus" - config = "config" +@cached +def get_blob_service(account_id: str) -> BlockBlobService: + logging.debug("getting blob container (account_id: %s)", account_id) + account_name, account_key = get_storage_account_name_key(account_id) + service = BlockBlobService(account_name=account_name, account_key=account_key) + return service -def get_account_id_by_type(storage_type: StorageType) -> str: - if storage_type == StorageType.corpus: - account_id = get_fuzz_storage() - elif storage_type == StorageType.config: - account_id = get_func_storage() - else: - raise NotImplementedError - return account_id - - -@cached(ttl=5) -def get_blob_service_by_type(storage_type: StorageType) -> Any: - account_id = get_account_id_by_type(storage_type) - return get_blob_service(account_id) +def get_service_by_container( + container: Container, storage_type: StorageType +) -> Optional[BlockBlobService]: + account = get_account_by_container(container, storage_type) + if account is None: + return None + service = get_blob_service(account) + return service -@cached(ttl=5) -def container_exists(name: str, storage_type: StorageType) -> bool: +def container_exists_on_account(container: Container, account_id: str) -> bool: try: - get_blob_service_by_type(storage_type).get_container_properties(name) + get_blob_service(account_id).get_container_properties(container) return True except AzureHttpError: return False -def get_containers(storage_type: StorageType) -> Dict[str, Dict[str, str]]: - return { - x.name: x.metadata - for x in get_blob_service_by_type(storage_type).list_containers( - include_metadata=True - ) - if not x.name.startswith("$") - } - - -def get_container_metadata( - name: str, storage_type: StorageType -) -> Optional[Dict[str, str]]: +def container_metadata(container: Container, account: str) -> Optional[Dict[str, str]]: try: - result = get_blob_service_by_type(storage_type).get_container_metadata(name) + result = get_blob_service(account).get_container_metadata(container) return cast(Dict[str, str], result) except AzureHttpError: pass return None -def create_container( - name: str, storage_type: StorageType, metadata: Optional[Dict[str, str]] +def get_account_by_container( + container: Container, storage_type: StorageType ) -> Optional[str]: - try: - get_blob_service_by_type(storage_type).create_container(name, metadata=metadata) - except AzureHttpError: - # azure storage already logs errors + accounts = get_accounts(storage_type) + + # check secondary accounts first by searching in reverse. + # + # By implementation, the primary account is specified first, followed by + # any secondary accounts. + # + # Secondary accounts, if they exist, are preferred for containers and have + # increased IOP rates, this should be a slight optimization + for account in reversed(accounts): + if container_exists_on_account(container, account): + return account + return None + + +def container_exists(container: Container, storage_type: StorageType) -> bool: + return get_account_by_container(container, storage_type) is not None + + +def get_containers(storage_type: StorageType) -> Dict[str, Dict[str, str]]: + containers: Dict[str, Dict[str, str]] = {} + + for account_id in get_accounts(storage_type): + containers.update( + { + x.name: x.metadata + for x in get_blob_service(account_id).list_containers( + include_metadata=True + ) + } + ) + + return containers + + +def get_container_metadata( + container: Container, storage_type: StorageType +) -> Optional[Dict[str, str]]: + account = get_account_by_container(container, storage_type) + if account is None: return None - return get_container_sas_url( - name, - storage_type, + return container_metadata(container, account) + + +def create_container( + container: Container, + storage_type: StorageType, + metadata: Optional[Dict[str, str]], +) -> Optional[str]: + service = get_service_by_container(container, storage_type) + if service is None: + account = choose_account(storage_type) + service = get_blob_service(account) + try: + service.create_container(container, metadata=metadata) + except AzureHttpError as err: + logging.error( + ( + "unable to create container. account: %s " + "container: %s metadata: %s - %s" + ), + account, + container, + metadata, + err, + ) + return None + + return get_container_sas_url_service( + container, + service, read=True, add=True, create=True, @@ -88,17 +141,19 @@ def create_container( ) -def delete_container(name: str, storage_type: StorageType) -> bool: - try: - return bool(get_blob_service_by_type(storage_type).delete_container(name)) - except AzureHttpError: - # azure storage already logs errors - return False +def delete_container(container: Container, storage_type: StorageType) -> bool: + accounts = get_accounts(storage_type) + for account in accounts: + service = get_blob_service(account) + if bool(service.delete_container(container)): + return True + return False -def get_container_sas_url( - container: str, - storage_type: StorageType, + +def get_container_sas_url_service( + container: Container, + service: BlockBlobService, *, read: bool = False, add: bool = False, @@ -107,7 +162,6 @@ def get_container_sas_url( delete: bool = False, list: bool = False, ) -> str: - service = get_blob_service_by_type(storage_type) expiry = datetime.datetime.utcnow() + datetime.timedelta(days=30) permission = ContainerPermissions(read, add, create, write, delete, list) @@ -120,8 +174,35 @@ def get_container_sas_url( return str(url) +def get_container_sas_url( + container: Container, + storage_type: StorageType, + *, + read: bool = False, + add: bool = False, + create: bool = False, + write: bool = False, + delete: bool = False, + list: bool = False, +) -> str: + service = get_service_by_container(container, storage_type) + if not service: + raise Exception("unable to create container sas for missing container") + + return get_container_sas_url_service( + container, + service, + read=read, + add=add, + create=create, + write=write, + delete=delete, + list=list, + ) + + def get_file_sas_url( - container: str, + container: Container, name: str, storage_type: StorageType, *, @@ -135,7 +216,10 @@ def get_file_sas_url( hours: int = 0, minutes: int = 0, ) -> str: - service = get_blob_service_by_type(storage_type) + service = get_service_by_container(container, storage_type) + if not service: + raise Exception("unable to find container: %s - %s" % (container, storage_type)) + expiry = datetime.datetime.utcnow() + datetime.timedelta( days=days, hours=hours, minutes=minutes ) @@ -150,18 +234,28 @@ def get_file_sas_url( def save_blob( - container: str, name: str, data: Union[str, bytes], storage_type: StorageType + container: Container, + name: str, + data: Union[str, bytes], + storage_type: StorageType, ) -> None: - service = get_blob_service_by_type(storage_type) - service.create_container(container) + service = get_service_by_container(container, storage_type) + if not service: + raise Exception("unable to find container: %s - %s" % (container, storage_type)) + if isinstance(data, str): service.create_blob_from_text(container, name, data) elif isinstance(data, bytes): service.create_blob_from_bytes(container, name, data) -def get_blob(container: str, name: str, storage_type: StorageType) -> Optional[bytes]: - service = get_blob_service_by_type(storage_type) +def get_blob( + container: Container, name: str, storage_type: StorageType +) -> Optional[bytes]: + service = get_service_by_container(container, storage_type) + if not service: + return None + try: blob = service.get_blob_to_bytes(container, name).content return cast(bytes, blob) @@ -169,8 +263,11 @@ def get_blob(container: str, name: str, storage_type: StorageType) -> Optional[b return None -def blob_exists(container: str, name: str, storage_type: StorageType) -> bool: - service = get_blob_service_by_type(storage_type) +def blob_exists(container: Container, name: str, storage_type: StorageType) -> bool: + service = get_service_by_container(container, storage_type) + if not service: + return False + try: service.get_blob_properties(container, name) return True @@ -178,8 +275,11 @@ def blob_exists(container: str, name: str, storage_type: StorageType) -> bool: return False -def delete_blob(container: str, name: str, storage_type: StorageType) -> bool: - service = get_blob_service_by_type(storage_type) +def delete_blob(container: Container, name: str, storage_type: StorageType) -> bool: + service = get_service_by_container(container, storage_type) + if not service: + return False + try: service.delete_blob(container, name) return True @@ -187,7 +287,7 @@ def delete_blob(container: str, name: str, storage_type: StorageType) -> bool: return False -def auth_download_url(container: str, filename: str) -> str: +def auth_download_url(container: Container, filename: str) -> str: instance = os.environ["ONEFUZZ_INSTANCE"] return "%s/api/download?%s" % ( instance, diff --git a/src/api-service/__app__/onefuzzlib/azure/creds.py b/src/api-service/__app__/onefuzzlib/azure/creds.py index d19b2d10d8..220babcd86 100644 --- a/src/api-service/__app__/onefuzzlib/azure/creds.py +++ b/src/api-service/__app__/onefuzzlib/azure/creds.py @@ -3,9 +3,8 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -import logging import os -from typing import Any, List, Optional, Tuple +from typing import Any, List from uuid import UUID from azure.cli.core import CLIError @@ -13,12 +12,11 @@ from azure.graphrbac import GraphRbacManagementClient from azure.graphrbac.models import CheckGroupMembershipParameters from azure.mgmt.resource import ResourceManagementClient -from azure.mgmt.storage import StorageManagementClient from azure.mgmt.subscription import SubscriptionClient -from azure.storage.blob import BlockBlobService from memoization import cached from msrestazure.azure_active_directory import MSIAuthentication from msrestazure.tools import parse_resource_id +from onefuzztypes.primitives import Container from .monkeypatch import allow_more_workers, reduce_logging @@ -35,34 +33,14 @@ def mgmt_client_factory(client_class: Any) -> Any: try: return get_client_from_cli_profile(client_class) except CLIError: - if issubclass(client_class, SubscriptionClient): - return client_class(get_msi()) - else: - return client_class(get_msi(), get_subscription()) + pass + except OSError: + pass - -@cached -def get_storage_account_name_key(account_id: Optional[str] = None) -> Tuple[str, str]: - db_client = mgmt_client_factory(StorageManagementClient) - if account_id is None: - account_id = os.environ["ONEFUZZ_DATA_STORAGE"] - resource = parse_resource_id(account_id) - key = ( - db_client.storage_accounts.list_keys( - resource["resource_group"], resource["name"] - ) - .keys[0] - .value - ) - return resource["name"], key - - -@cached -def get_blob_service(account_id: Optional[str] = None) -> BlockBlobService: - logging.debug("getting blob container (account_id: %s)", account_id) - name, key = get_storage_account_name_key(account_id) - service = BlockBlobService(account_name=name, account_key=key) - return service + if issubclass(client_class, SubscriptionClient): + return client_class(get_msi()) + else: + return client_class(get_msi(), get_subscription()) @cached @@ -92,16 +70,6 @@ def get_insights_appid() -> str: return os.environ["APPINSIGHTS_APPID"] -# @cached -def get_fuzz_storage() -> str: - return os.environ["ONEFUZZ_DATA_STORAGE"] - - -# @cached -def get_func_storage() -> str: - return os.environ["ONEFUZZ_FUNC_STORAGE"] - - @cached def get_instance_name() -> str: return os.environ["ONEFUZZ_INSTANCE_NAME"] @@ -114,9 +82,10 @@ def get_instance_url() -> str: @cached def get_instance_id() -> UUID: - from .containers import StorageType, get_blob + from .containers import get_blob + from .storage import StorageType - blob = get_blob("base-config", "instance_id", StorageType.config) + blob = get_blob(Container("base-config"), "instance_id", StorageType.config) if blob is None: raise Exception("missing instance_id") return UUID(blob.decode()) diff --git a/src/api-service/__app__/onefuzzlib/azure/queue.py b/src/api-service/__app__/onefuzzlib/azure/queue.py index e73edf35b1..37c28303ce 100644 --- a/src/api-service/__app__/onefuzzlib/azure/queue.py +++ b/src/api-service/__app__/onefuzzlib/azure/queue.py @@ -19,8 +19,7 @@ from memoization import cached from pydantic import BaseModel -from .containers import StorageType, get_account_id_by_type -from .creds import get_storage_account_name_key +from .storage import StorageType, get_primary_account, get_storage_account_name_key QueueNameType = Union[str, UUID] @@ -29,7 +28,7 @@ @cached(ttl=60) def get_queue_client(storage_type: StorageType) -> QueueServiceClient: - account_id = get_account_id_by_type(storage_type) + account_id = get_primary_account(storage_type) logging.debug("getting blob container (account_id: %s)", account_id) name, key = get_storage_account_name_key(account_id) account_url = "https://%s.queue.core.windows.net" % name @@ -50,7 +49,7 @@ def get_queue_sas( update: bool = False, process: bool = False, ) -> str: - account_id = get_account_id_by_type(storage_type) + account_id = get_primary_account(storage_type) logging.debug("getting queue sas %s (account_id: %s)", queue, account_id) name, key = get_storage_account_name_key(account_id) expiry = datetime.datetime.utcnow() + datetime.timedelta(days=30) diff --git a/src/api-service/__app__/onefuzzlib/azure/storage.py b/src/api-service/__app__/onefuzzlib/azure/storage.py new file mode 100644 index 0000000000..460d7b2cc9 --- /dev/null +++ b/src/api-service/__app__/onefuzzlib/azure/storage.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import logging +import os +import random +from enum import Enum +from typing import List, Tuple + +from azure.mgmt.storage import StorageManagementClient +from memoization import cached +from msrestazure.tools import parse_resource_id + +from .creds import get_base_resource_group, mgmt_client_factory + + +class StorageType(Enum): + corpus = "corpus" + config = "config" + + +@cached +def get_fuzz_storage() -> str: + return os.environ["ONEFUZZ_DATA_STORAGE"] + + +@cached +def get_func_storage() -> str: + return os.environ["ONEFUZZ_FUNC_STORAGE"] + + +@cached +def get_primary_account(storage_type: StorageType) -> str: + if storage_type == StorageType.corpus: + # see #322 for discussion about typing + return get_fuzz_storage() + elif storage_type == StorageType.config: + # see #322 for discussion about typing + return get_func_storage() + raise NotImplementedError + + +@cached +def get_accounts(storage_type: StorageType) -> List[str]: + if storage_type == StorageType.corpus: + return corpus_accounts() + elif storage_type == StorageType.config: + return [get_func_storage()] + else: + raise NotImplementedError + + +@cached +def get_storage_account_name_key(account_id: str) -> Tuple[str, str]: + client = mgmt_client_factory(StorageManagementClient) + resource = parse_resource_id(account_id) + key = ( + client.storage_accounts.list_keys(resource["resource_group"], resource["name"]) + .keys[0] + .value + ) + return resource["name"], key + + +def choose_account(storage_type: StorageType) -> str: + accounts = get_accounts(storage_type) + if not accounts: + raise Exception(f"no storage accounts for {storage_type}") + + if len(accounts) == 1: + return accounts[0] + + # Use a random secondary storage account if any are available. This + # reduces IOP contention for the Storage Queues, which are only available + # on primary accounts + return random.choice(accounts[1:]) + + +@cached +def corpus_accounts() -> List[str]: + skip = get_func_storage() + results = [get_fuzz_storage()] + + client = mgmt_client_factory(StorageManagementClient) + group = get_base_resource_group() + for account in client.storage_accounts.list_by_resource_group(group): + # protection from someone adding the corpus tag to the config account + if account.id == skip: + continue + + if account.id in results: + continue + + if account.primary_endpoints.blob is None: + continue + + if ( + "storage_type" not in account.tags + or account.tags["storage_type"] != "corpus" + ): + continue + + results.append(account.id) + + logging.info("corpus accounts: %s", corpus_accounts) + return results diff --git a/src/api-service/__app__/onefuzzlib/azure/table.py b/src/api-service/__app__/onefuzzlib/azure/table.py index 1b3df2e05b..8272ad5c5f 100644 --- a/src/api-service/__app__/onefuzzlib/azure/table.py +++ b/src/api-service/__app__/onefuzzlib/azure/table.py @@ -10,7 +10,7 @@ from azure.cosmosdb.table import TableService from memoization import cached -from .creds import get_storage_account_name_key +from .storage import get_storage_account_name_key @cached(ttl=60) diff --git a/src/api-service/__app__/onefuzzlib/extension.py b/src/api-service/__app__/onefuzzlib/extension.py index 2e5e5c8afd..b001421154 100644 --- a/src/api-service/__app__/onefuzzlib/extension.py +++ b/src/api-service/__app__/onefuzzlib/extension.py @@ -9,17 +9,13 @@ from onefuzztypes.enums import OS, AgentMode from onefuzztypes.models import AgentConfig, Pool, ReproConfig, Scaleset -from onefuzztypes.primitives import Extension, Region - -from .azure.containers import ( - StorageType, - get_container_sas_url, - get_file_sas_url, - save_blob, -) +from onefuzztypes.primitives import Container, Extension, Region + +from .azure.containers import get_container_sas_url, get_file_sas_url, save_blob from .azure.creds import get_instance_id, get_instance_url from .azure.monitor import get_monitor_settings from .azure.queue import get_queue_sas +from .azure.storage import StorageType from .reports import get_report @@ -95,8 +91,12 @@ def build_scaleset_script(pool: Pool, scaleset: Scaleset) -> str: ssh_path = "$env:ProgramData/ssh/administrators_authorized_keys" commands += [f'Set-Content -Path {ssh_path} -Value "{ssh_key}"'] - save_blob("vm-scripts", filename, sep.join(commands) + sep, StorageType.config) - return get_file_sas_url("vm-scripts", filename, StorageType.config, read=True) + save_blob( + Container("vm-scripts"), filename, sep.join(commands) + sep, StorageType.config + ) + return get_file_sas_url( + Container("vm-scripts"), filename, StorageType.config, read=True + ) def build_pool_config(pool: Pool) -> str: @@ -116,14 +116,14 @@ def build_pool_config(pool: Pool) -> str: filename = f"{pool.name}/config.json" save_blob( - "vm-scripts", + Container("vm-scripts"), filename, config.json(), StorageType.config, ) return get_file_sas_url( - "vm-scripts", + Container("vm-scripts"), filename, StorageType.config, read=True, @@ -135,24 +135,28 @@ def update_managed_scripts() -> None: "azcopy sync '%s' instance-specific-setup" % ( get_container_sas_url( - "instance-specific-setup", + Container("instance-specific-setup"), StorageType.config, read=True, list=True, ) ), "azcopy sync '%s' tools" - % (get_container_sas_url("tools", StorageType.config, read=True, list=True)), + % ( + get_container_sas_url( + Container("tools"), StorageType.config, read=True, list=True + ) + ), ] save_blob( - "vm-scripts", + Container("vm-scripts"), "managed.ps1", "\r\n".join(commands) + "\r\n", StorageType.config, ) save_blob( - "vm-scripts", + Container("vm-scripts"), "managed.sh", "\n".join(commands) + "\n", StorageType.config, @@ -170,25 +174,25 @@ def agent_config( if vm_os == OS.windows: urls += [ get_file_sas_url( - "vm-scripts", + Container("vm-scripts"), "managed.ps1", StorageType.config, read=True, ), get_file_sas_url( - "tools", + Container("tools"), "win64/azcopy.exe", StorageType.config, read=True, ), get_file_sas_url( - "tools", + Container("tools"), "win64/setup.ps1", StorageType.config, read=True, ), get_file_sas_url( - "tools", + Container("tools"), "win64/onefuzz.ps1", StorageType.config, read=True, @@ -212,19 +216,19 @@ def agent_config( elif vm_os == OS.linux: urls += [ get_file_sas_url( - "vm-scripts", + Container("vm-scripts"), "managed.sh", StorageType.config, read=True, ), get_file_sas_url( - "tools", + Container("tools"), "linux/azcopy", StorageType.config, read=True, ), get_file_sas_url( - "tools", + Container("tools"), "linux/setup.sh", StorageType.config, read=True, @@ -260,7 +264,7 @@ def repro_extensions( repro_os: OS, repro_id: UUID, repro_config: ReproConfig, - setup_container: Optional[str], + setup_container: Optional[Container], ) -> List[Extension]: # TODO - what about contents of repro.ps1 / repro.sh? report = get_report(repro_config.container, repro_config.path) @@ -302,7 +306,7 @@ def repro_extensions( script_name = "task-setup.sh" save_blob( - "task-configs", + Container("task-configs"), "%s/%s" % (repro_id, script_name), task_script, StorageType.config, @@ -311,13 +315,13 @@ def repro_extensions( for repro_file in repro_files: urls += [ get_file_sas_url( - "repro-scripts", + Container("repro-scripts"), repro_file, StorageType.config, read=True, ), get_file_sas_url( - "task-configs", + Container("task-configs"), "%s/%s" % (repro_id, script_name), StorageType.config, read=True, @@ -333,13 +337,13 @@ def repro_extensions( def proxy_manager_extensions(region: Region) -> List[Extension]: urls = [ get_file_sas_url( - "proxy-configs", + Container("proxy-configs"), "%s/config.json" % region, StorageType.config, read=True, ), get_file_sas_url( - "tools", + Container("tools"), "linux/onefuzz-proxy-manager", StorageType.config, read=True, diff --git a/src/api-service/__app__/onefuzzlib/notifications/ado.py b/src/api-service/__app__/onefuzzlib/notifications/ado.py index bb433467da..963f6fe868 100644 --- a/src/api-service/__app__/onefuzzlib/notifications/ado.py +++ b/src/api-service/__app__/onefuzzlib/notifications/ado.py @@ -25,6 +25,7 @@ ) from memoization import cached from onefuzztypes.models import ADOTemplate, Report +from onefuzztypes.primitives import Container from .common import Render, fail_task @@ -49,7 +50,7 @@ def get_valid_fields( class ADO: def __init__( - self, container: str, filename: str, config: ADOTemplate, report: Report + self, container: Container, filename: str, config: ADOTemplate, report: Report ): self.config = config self.renderer = Render(container, filename, report) @@ -200,7 +201,7 @@ def process(self) -> None: def notify_ado( - config: ADOTemplate, container: str, filename: str, report: Report + config: ADOTemplate, container: Container, filename: str, report: Report ) -> None: logging.info( "notify ado: job_id:%s task_id:%s container:%s filename:%s", diff --git a/src/api-service/__app__/onefuzzlib/notifications/common.py b/src/api-service/__app__/onefuzzlib/notifications/common.py index eca196473f..4c30f8aff7 100644 --- a/src/api-service/__app__/onefuzzlib/notifications/common.py +++ b/src/api-service/__app__/onefuzzlib/notifications/common.py @@ -9,6 +9,7 @@ from jinja2.sandbox import SandboxedEnvironment from onefuzztypes.enums import ErrorCode from onefuzztypes.models import Error, Report +from onefuzztypes.primitives import Container from ..azure.containers import auth_download_url from ..azure.creds import get_instance_url @@ -33,7 +34,7 @@ def fail_task(report: Report, error: Exception) -> None: class Render: - def __init__(self, container: str, filename: str, report: Report): + def __init__(self, container: Container, filename: str, report: Report): self.report = report self.container = container self.filename = filename diff --git a/src/api-service/__app__/onefuzzlib/notifications/github_issues.py b/src/api-service/__app__/onefuzzlib/notifications/github_issues.py index 452e7b3574..c3fa7364e9 100644 --- a/src/api-service/__app__/onefuzzlib/notifications/github_issues.py +++ b/src/api-service/__app__/onefuzzlib/notifications/github_issues.py @@ -11,13 +11,18 @@ from github3.issues import Issue from onefuzztypes.enums import GithubIssueSearchMatch from onefuzztypes.models import GithubIssueTemplate, Report +from onefuzztypes.primitives import Container from .common import Render, fail_task class GithubIssue: def __init__( - self, config: GithubIssueTemplate, container: str, filename: str, report: Report + self, + config: GithubIssueTemplate, + container: Container, + filename: str, + report: Report, ): self.config = config self.report = report @@ -95,7 +100,10 @@ def process(self) -> None: def github_issue( - config: GithubIssueTemplate, container: str, filename: str, report: Optional[Report] + config: GithubIssueTemplate, + container: Container, + filename: str, + report: Optional[Report], ) -> None: if report is None: return diff --git a/src/api-service/__app__/onefuzzlib/notifications/main.py b/src/api-service/__app__/onefuzzlib/notifications/main.py index fdcbc0c2b3..ad798031fb 100644 --- a/src/api-service/__app__/onefuzzlib/notifications/main.py +++ b/src/api-service/__app__/onefuzzlib/notifications/main.py @@ -21,12 +21,12 @@ from onefuzztypes.primitives import Container, Event from ..azure.containers import ( - StorageType, container_exists, get_container_metadata, get_file_sas_url, ) from ..azure.queue import send_message +from ..azure.storage import StorageType from ..dashboard import add_event from ..orm import ORMMixin from ..reports import get_report diff --git a/src/api-service/__app__/onefuzzlib/notifications/teams.py b/src/api-service/__app__/onefuzzlib/notifications/teams.py index 7389f0eeec..9d501cbea9 100644 --- a/src/api-service/__app__/onefuzzlib/notifications/teams.py +++ b/src/api-service/__app__/onefuzzlib/notifications/teams.py @@ -8,6 +8,7 @@ import requests from onefuzztypes.models import Report, TeamsTemplate +from onefuzztypes.primitives import Container from ..azure.containers import auth_download_url from ..tasks.config import get_setup_container @@ -51,7 +52,7 @@ def send_teams_webhook( def notify_teams( - config: TeamsTemplate, container: str, filename: str, report: Optional[Report] + config: TeamsTemplate, container: Container, filename: str, report: Optional[Report] ) -> None: text = None facts: List[Dict[str, str]] = [] diff --git a/src/api-service/__app__/onefuzzlib/pools.py b/src/api-service/__app__/onefuzzlib/pools.py index c184f01985..79b0f54644 100644 --- a/src/api-service/__app__/onefuzzlib/pools.py +++ b/src/api-service/__app__/onefuzzlib/pools.py @@ -36,7 +36,6 @@ from .__version__ import __version__ from .azure.auth import build_auth -from .azure.containers import StorageType from .azure.image import get_os from .azure.network import Network from .azure.queue import ( @@ -47,6 +46,7 @@ queue_object, remove_first_message, ) +from .azure.storage import StorageType from .azure.vmss import ( UnableToUpdate, create_vmss, diff --git a/src/api-service/__app__/onefuzzlib/proxy.py b/src/api-service/__app__/onefuzzlib/proxy.py index 64ddae8f5f..edb07c33a4 100644 --- a/src/api-service/__app__/onefuzzlib/proxy.py +++ b/src/api-service/__app__/onefuzzlib/proxy.py @@ -16,14 +16,15 @@ ProxyConfig, ProxyHeartbeat, ) -from onefuzztypes.primitives import Region +from onefuzztypes.primitives import Container, Region from pydantic import Field from .__version__ import __version__ from .azure.auth import build_auth -from .azure.containers import StorageType, get_file_sas_url, save_blob +from .azure.containers import get_file_sas_url, save_blob from .azure.ip import get_public_ip from .azure.queue import get_queue_sas +from .azure.storage import StorageType from .azure.vm import VM from .extension import proxy_manager_extensions from .orm import MappingIntStrAny, ORMMixin, QueryFilter @@ -188,7 +189,7 @@ def save_proxy_config(self) -> None: forwards = self.get_forwards() proxy_config = ProxyConfig( url=get_file_sas_url( - "proxy-configs", + Container("proxy-configs"), "%s/config.json" % self.region, StorageType.config, read=True, @@ -203,7 +204,7 @@ def save_proxy_config(self) -> None: ) save_blob( - "proxy-configs", + Container("proxy-configs"), "%s/config.json" % self.region, proxy_config.json(), StorageType.config, diff --git a/src/api-service/__app__/onefuzzlib/reports.py b/src/api-service/__app__/onefuzzlib/reports.py index 7255b25914..36442f332c 100644 --- a/src/api-service/__app__/onefuzzlib/reports.py +++ b/src/api-service/__app__/onefuzzlib/reports.py @@ -8,9 +8,11 @@ from typing import Optional, Union from onefuzztypes.models import Report +from onefuzztypes.primitives import Container from pydantic import ValidationError -from .azure.containers import StorageType, get_blob +from .azure.containers import get_blob +from .azure.storage import StorageType def parse_report( @@ -44,7 +46,7 @@ def parse_report( return entry -def get_report(container: str, filename: str) -> Optional[Report]: +def get_report(container: Container, filename: str) -> Optional[Report]: metadata = "/".join([container, filename]) if not filename.endswith(".json"): logging.error("get_report invalid extension: %s", metadata) diff --git a/src/api-service/__app__/onefuzzlib/repro.py b/src/api-service/__app__/onefuzzlib/repro.py index 929dbf7bdd..e837426c6a 100644 --- a/src/api-service/__app__/onefuzzlib/repro.py +++ b/src/api-service/__app__/onefuzzlib/repro.py @@ -12,11 +12,13 @@ from onefuzztypes.models import Error from onefuzztypes.models import Repro as BASE_REPRO from onefuzztypes.models import ReproConfig, TaskVm, UserInfo +from onefuzztypes.primitives import Container from .azure.auth import build_auth -from .azure.containers import StorageType, save_blob +from .azure.containers import save_blob from .azure.creds import get_base_region from .azure.ip import get_public_ip +from .azure.storage import StorageType from .azure.vm import VM from .extension import repro_extensions from .orm import ORMMixin, QueryFilter @@ -98,7 +100,7 @@ def set_failed(self, vm_data: VirtualMachine) -> None: ) return self.set_error(Error(code=ErrorCode.VM_CREATE_FAILED, errors=errors)) - def get_setup_container(self) -> Optional[str]: + def get_setup_container(self) -> Optional[Container]: task = Task.get_by_task_id(self.task_id) if isinstance(task, Task): for container in task.config.containers: @@ -202,7 +204,7 @@ def build_repro_script(self) -> Optional[Error]: for filename in files: save_blob( - "repro-scripts", + Container("repro-scripts"), "%s/%s" % (self.vm_id, filename), files[filename], StorageType.config, diff --git a/src/api-service/__app__/onefuzzlib/tasks/config.py b/src/api-service/__app__/onefuzzlib/tasks/config.py index c913a6c3f3..3df5f9cb6c 100644 --- a/src/api-service/__app__/onefuzzlib/tasks/config.py +++ b/src/api-service/__app__/onefuzzlib/tasks/config.py @@ -10,15 +10,12 @@ from onefuzztypes.enums import Compare, ContainerPermission, ContainerType, TaskFeature from onefuzztypes.models import TaskConfig, TaskDefinition, TaskUnitConfig +from onefuzztypes.primitives import Container -from ..azure.containers import ( - StorageType, - blob_exists, - container_exists, - get_container_sas_url, -) +from ..azure.containers import blob_exists, container_exists, get_container_sas_url from ..azure.creds import get_instance_id, get_instance_url from ..azure.queue import get_queue_sas +from ..azure.storage import StorageType from .defs import TASK_DEFINITIONS LOGGER = logging.getLogger("onefuzz") @@ -334,7 +331,7 @@ def build_task_config( return config -def get_setup_container(config: TaskConfig) -> str: +def get_setup_container(config: TaskConfig) -> Container: for container in config.containers: if container.type == ContainerType.setup: return container.name diff --git a/src/api-service/__app__/onefuzzlib/tasks/main.py b/src/api-service/__app__/onefuzzlib/tasks/main.py index 64786398df..3827989cc5 100644 --- a/src/api-service/__app__/onefuzzlib/tasks/main.py +++ b/src/api-service/__app__/onefuzzlib/tasks/main.py @@ -18,9 +18,9 @@ WebhookEventTaskStopped, ) -from ..azure.containers import StorageType from ..azure.image import get_os from ..azure.queue import create_queue, delete_queue +from ..azure.storage import StorageType from ..orm import MappingIntStrAny, ORMMixin, QueryFilter from ..pools import Node, Pool, Scaleset from ..proxy_forward import ProxyForward diff --git a/src/api-service/__app__/onefuzzlib/tasks/scheduler.py b/src/api-service/__app__/onefuzzlib/tasks/scheduler.py index c96c6ffbcc..a24265a54a 100644 --- a/src/api-service/__app__/onefuzzlib/tasks/scheduler.py +++ b/src/api-service/__app__/onefuzzlib/tasks/scheduler.py @@ -11,7 +11,8 @@ from onefuzztypes.models import WorkSet, WorkUnit from pydantic import BaseModel -from ..azure.containers import StorageType, blob_exists, get_container_sas_url +from ..azure.containers import blob_exists, get_container_sas_url +from ..azure.storage import StorageType from ..pools import Pool from .config import build_task_config, get_setup_container from .main import Task diff --git a/src/api-service/__app__/onefuzzlib/updates.py b/src/api-service/__app__/onefuzzlib/updates.py index ab74872c34..c4a842e79e 100644 --- a/src/api-service/__app__/onefuzzlib/updates.py +++ b/src/api-service/__app__/onefuzzlib/updates.py @@ -10,8 +10,8 @@ from onefuzztypes.enums import UpdateType from pydantic import BaseModel -from .azure.containers import StorageType from .azure.queue import queue_object +from .azure.storage import StorageType # This class isn't intended to be shared outside of the service diff --git a/src/api-service/__app__/onefuzzlib/webhooks.py b/src/api-service/__app__/onefuzzlib/webhooks.py index 3b7102b44b..02b8d62c3c 100644 --- a/src/api-service/__app__/onefuzzlib/webhooks.py +++ b/src/api-service/__app__/onefuzzlib/webhooks.py @@ -27,8 +27,8 @@ from pydantic import BaseModel from .__version__ import __version__ -from .azure.containers import StorageType from .azure.queue import queue_object +from .azure.storage import StorageType from .orm import ORMMixin MAX_TRIES = 5 diff --git a/src/api-service/__app__/pool/__init__.py b/src/api-service/__app__/pool/__init__.py index cae14eae66..68359a8eb5 100644 --- a/src/api-service/__app__/pool/__init__.py +++ b/src/api-service/__app__/pool/__init__.py @@ -12,7 +12,6 @@ from onefuzztypes.requests import PoolCreate, PoolSearch, PoolStop from onefuzztypes.responses import BoolResult -from ..onefuzzlib.azure.containers import StorageType from ..onefuzzlib.azure.creds import ( get_base_region, get_instance_id, @@ -20,6 +19,7 @@ get_regions, ) from ..onefuzzlib.azure.queue import get_queue_sas +from ..onefuzzlib.azure.storage import StorageType from ..onefuzzlib.azure.vmss import list_available_skus from ..onefuzzlib.endpoint_authorization import call_if_user from ..onefuzzlib.pools import Pool diff --git a/src/api-service/__app__/queue_file_changes/__init__.py b/src/api-service/__app__/queue_file_changes/__init__.py index 5101debedf..3aa4b83719 100644 --- a/src/api-service/__app__/queue_file_changes/__init__.py +++ b/src/api-service/__app__/queue_file_changes/__init__.py @@ -9,7 +9,7 @@ import azure.functions as func -from ..onefuzzlib.azure.creds import get_fuzz_storage +from ..onefuzzlib.azure.storage import corpus_accounts from ..onefuzzlib.dashboard import get_event from ..onefuzzlib.notifications.main import new_files @@ -25,7 +25,7 @@ def file_added(event: Dict) -> None: def main(msg: func.QueueMessage, dashboard: func.Out[str]) -> None: event = json.loads(msg.get_body()) - if event["topic"] != get_fuzz_storage(): + if event["topic"] in corpus_accounts(): return if event["eventType"] != "Microsoft.Storage.BlobCreated": diff --git a/src/deployment/azuredeploy.json b/src/deployment/azuredeploy.json index 2ea7b8f476..614f8b4a4d 100644 --- a/src/deployment/azuredeploy.json +++ b/src/deployment/azuredeploy.json @@ -580,6 +580,14 @@ "[resourceId('Microsoft.Storage/storageAccounts', variables('storageAccountName'))]" ] }, + { + "type": "Microsoft.Storage/storageAccounts/blobServices/containers", + "apiVersion": "2018-03-01-preview", + "name": "[concat(variables('storageAccountNameFunc'), '/default/', 'vm-scripts')]", + "dependsOn": [ + "[resourceId('Microsoft.Storage/storageAccounts', variables('storageAccountNameFunc'))]" + ] + }, { "type": "Microsoft.Storage/storageAccounts/blobServices/containers", "apiVersion": "2018-03-01-preview", diff --git a/src/utils/add-corpus-storage-accounts/add-corpus-storage-account.py b/src/utils/add-corpus-storage-accounts/add-corpus-storage-account.py new file mode 100644 index 0000000000..0f0497fb50 --- /dev/null +++ b/src/utils/add-corpus-storage-accounts/add-corpus-storage-account.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python +# +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +import argparse +import json +import uuid + +from azure.common.client_factory import get_client_from_cli_profile +from azure.mgmt.eventgrid import EventGridManagementClient +from azure.mgmt.eventgrid.models import EventSubscription +from azure.mgmt.storage import StorageManagementClient +from azure.mgmt.storage.models import ( + AccessTier, + Kind, + Sku, + SkuName, + StorageAccountCreateParameters, +) + +# This was generated randomly and should be preserved moving forwards +STORAGE_GUID_NAMESPACE = uuid.UUID("f7eb528c-d849-4b81-9046-e7036f6203df") + + +def get_base_event( + client: EventGridManagementClient, resource_group: str, location: str +) -> EventSubscription: + for entry in client.event_subscriptions.list_regional_by_resource_group( + resource_group, location + ): + if ( + entry.name == "onefuzz1" + and entry.type == "Microsoft.EventGrid/eventSubscriptions" + and entry.event_delivery_schema == "EventGridSchema" + and entry.destination.endpoint_type == "StorageQueue" + and entry.destination.queue_name == "file-changes" + ): + return entry + + raise Exception("unable to find base eventgrid subscription") + + +def add_event_grid(src_account_id: str, resource_group: str, location: str) -> None: + client = get_client_from_cli_profile(EventGridManagementClient) + base = get_base_event(client, resource_group, location) + + event_subscription_info = EventSubscription( + destination=base.destination, + filter=base.filter, + retry_policy=base.retry_policy, + ) + + topic_id = uuid.uuid5(STORAGE_GUID_NAMESPACE, src_account_id).hex + + result = client.event_subscriptions.create_or_update( + src_account_id, "corpus" + topic_id, event_subscription_info + ).result() + if result.provisioning_state != "Succeeded": + raise Exception( + "eventgrid subscription failed: %s" + % json.dumps(result.as_dict(), indent=4, sort_keys=True), + ) + + +def create_storage(resource_group: str, account_name: str, location: str) -> str: + params = StorageAccountCreateParameters( + sku=Sku(name=SkuName.premium_lrs), + kind=Kind.block_blob_storage, + location=location, + tags={"storage_type": "corpus"}, + access_tier=AccessTier.hot, + allow_blob_public_access=False, + minimum_tls_version="TLS1_2", + ) + + client = get_client_from_cli_profile(StorageManagementClient) + account = client.storage_accounts.create( + resource_group, account_name, params + ).result() + if account.provisioning_state != "Succeeded": + raise Exception( + "storage account creation failed: %s", + json.dumps(account.as_dict(), indent=4, sort_keys=True), + ) + return account.id + + +def create(resource_group: str, account_name: str, location: str) -> None: + new_account_id = create_storage(resource_group, account_name, location) + add_event_grid(new_account_id, resource_group, location) + + +def main(): + formatter = argparse.ArgumentDefaultsHelpFormatter + parser = argparse.ArgumentParser(formatter_class=formatter) + parser.add_argument("resource_group") + parser.add_argument("account_name") + parser.add_argument("location") + args = parser.parse_args() + + create(args.resource_group, args.account_name, args.location) + + +if __name__ == "__main__": + main() diff --git a/src/utils/add-corpus-storage-accounts/requirements-lint.txt b/src/utils/add-corpus-storage-accounts/requirements-lint.txt new file mode 100644 index 0000000000..27a97b8b59 --- /dev/null +++ b/src/utils/add-corpus-storage-accounts/requirements-lint.txt @@ -0,0 +1,6 @@ +flake8 +mypy +pytest +isort +vulture +black diff --git a/src/utils/add-corpus-storage-accounts/requirements.txt b/src/utils/add-corpus-storage-accounts/requirements.txt new file mode 100644 index 0000000000..302e9cbe3c --- /dev/null +++ b/src/utils/add-corpus-storage-accounts/requirements.txt @@ -0,0 +1,3 @@ +azure-mgmt-storage~=11.2.0 +azure-cli-core==2.13.0 +azure-mgmt-eventgrid==3.0.0rc7 \ No newline at end of file