mirror of
https://github.com/microsoft/onefuzz.git
synced 2025-06-20 05:23:44 +00:00
Use Storage Account types, rather than account_id (#320)
We need to move to supporting data sharding. One of the steps towards that is stop passing around `account_id`, rather we need to specify the type of storage we need.
This commit is contained in:
@ -11,13 +11,13 @@ from uuid import UUID
|
||||
from onefuzztypes.enums import Compare, ContainerPermission, ContainerType, TaskFeature
|
||||
from onefuzztypes.models import TaskConfig, TaskDefinition, TaskUnitConfig
|
||||
|
||||
from ..azure.containers import blob_exists, container_exists, get_container_sas_url
|
||||
from ..azure.creds import (
|
||||
get_func_storage,
|
||||
get_fuzz_storage,
|
||||
get_instance_id,
|
||||
get_instance_url,
|
||||
from ..azure.containers import (
|
||||
StorageType,
|
||||
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 .defs import TASK_DEFINITIONS
|
||||
|
||||
@ -68,7 +68,7 @@ def check_containers(definition: TaskDefinition, config: TaskConfig) -> None:
|
||||
containers: Dict[ContainerType, List[str]] = {}
|
||||
for container in config.containers:
|
||||
if container.name not in checked:
|
||||
if not container_exists(container.name):
|
||||
if not container_exists(container.name, StorageType.corpus):
|
||||
raise TaskConfigError("missing container: %s" % container.name)
|
||||
checked.add(container.name)
|
||||
|
||||
@ -137,7 +137,7 @@ def check_config(config: TaskConfig) -> None:
|
||||
|
||||
if TaskFeature.target_exe in definition.features:
|
||||
container = [x for x in config.containers if x.type == ContainerType.setup][0]
|
||||
if not blob_exists(container.name, config.task.target_exe):
|
||||
if not blob_exists(container.name, config.task.target_exe, StorageType.corpus):
|
||||
err = "target_exe `%s` does not exist in the setup container `%s`" % (
|
||||
config.task.target_exe,
|
||||
container.name,
|
||||
@ -153,7 +153,7 @@ def check_config(config: TaskConfig) -> None:
|
||||
for tool_path in tools_paths:
|
||||
if config.task.generator_exe.startswith(tool_path):
|
||||
generator = config.task.generator_exe.replace(tool_path, "")
|
||||
if not blob_exists(container.name, generator):
|
||||
if not blob_exists(container.name, generator, StorageType.corpus):
|
||||
err = (
|
||||
"generator_exe `%s` does not exist in the tools container `%s`"
|
||||
% (
|
||||
@ -188,7 +188,7 @@ def build_task_config(
|
||||
telemetry_key=os.environ.get("ONEFUZZ_TELEMETRY"),
|
||||
heartbeat_queue=get_queue_sas(
|
||||
"task-heartbeat",
|
||||
account_id=get_func_storage(),
|
||||
StorageType.config,
|
||||
add=True,
|
||||
),
|
||||
back_channel_address="https://%s/api/back_channel" % (get_instance_url()),
|
||||
@ -198,11 +198,11 @@ def build_task_config(
|
||||
if definition.monitor_queue:
|
||||
config.input_queue = get_queue_sas(
|
||||
task_id,
|
||||
StorageType.corpus,
|
||||
add=True,
|
||||
read=True,
|
||||
update=True,
|
||||
process=True,
|
||||
account_id=get_fuzz_storage(),
|
||||
)
|
||||
|
||||
for container_def in definition.containers:
|
||||
@ -219,6 +219,7 @@ def build_task_config(
|
||||
"path": "_".join(["task", container_def.type.name, str(i)]),
|
||||
"url": get_container_sas_url(
|
||||
container.name,
|
||||
StorageType.corpus,
|
||||
read=ContainerPermission.Read in container_def.permissions,
|
||||
write=ContainerPermission.Write in container_def.permissions,
|
||||
add=ContainerPermission.Add in container_def.permissions,
|
||||
|
@ -18,7 +18,7 @@ from onefuzztypes.webhooks import (
|
||||
WebhookEventTaskStopped,
|
||||
)
|
||||
|
||||
from ..azure.creds import get_fuzz_storage
|
||||
from ..azure.containers import StorageType
|
||||
from ..azure.image import get_os
|
||||
from ..azure.queue import create_queue, delete_queue
|
||||
from ..orm import MappingIntStrAny, ORMMixin, QueryFilter
|
||||
@ -123,7 +123,7 @@ class Task(BASE_TASK, ORMMixin):
|
||||
}
|
||||
|
||||
def init(self) -> None:
|
||||
create_queue(self.task_id, account_id=get_fuzz_storage())
|
||||
create_queue(self.task_id, StorageType.corpus)
|
||||
self.state = TaskState.waiting
|
||||
self.save()
|
||||
|
||||
@ -132,7 +132,7 @@ class Task(BASE_TASK, ORMMixin):
|
||||
|
||||
logging.info("stopping task: %s:%s", self.job_id, self.task_id)
|
||||
ProxyForward.remove_forward(self.task_id)
|
||||
delete_queue(str(self.task_id), account_id=get_fuzz_storage())
|
||||
delete_queue(str(self.task_id), StorageType.corpus)
|
||||
Node.stop_task(self.task_id)
|
||||
self.state = TaskState.stopped
|
||||
self.save()
|
||||
|
@ -10,8 +10,12 @@ from uuid import UUID
|
||||
from onefuzztypes.enums import OS, PoolState, TaskState
|
||||
from onefuzztypes.models import WorkSet, WorkUnit
|
||||
|
||||
from ..azure.containers import blob_exists, get_container_sas_url, save_blob
|
||||
from ..azure.creds import get_func_storage
|
||||
from ..azure.containers import (
|
||||
StorageType,
|
||||
blob_exists,
|
||||
get_container_sas_url,
|
||||
save_blob,
|
||||
)
|
||||
from ..pools import Pool
|
||||
from .config import build_task_config, get_setup_container
|
||||
from .main import Task
|
||||
@ -60,20 +64,26 @@ def schedule_tasks() -> None:
|
||||
agent_config = build_task_config(task.job_id, task.task_id, task.config)
|
||||
|
||||
setup_container = get_setup_container(task.config)
|
||||
setup_url = get_container_sas_url(setup_container, read=True, list=True)
|
||||
setup_url = get_container_sas_url(
|
||||
setup_container, StorageType.corpus, read=True, list=True
|
||||
)
|
||||
|
||||
setup_script = None
|
||||
|
||||
if task.os == OS.windows and blob_exists(setup_container, "setup.ps1"):
|
||||
if task.os == OS.windows and blob_exists(
|
||||
setup_container, "setup.ps1", StorageType.corpus
|
||||
):
|
||||
setup_script = "setup.ps1"
|
||||
if task.os == OS.linux and blob_exists(setup_container, "setup.sh"):
|
||||
if task.os == OS.linux and blob_exists(
|
||||
setup_container, "setup.sh", StorageType.corpus
|
||||
):
|
||||
setup_script = "setup.sh"
|
||||
|
||||
save_blob(
|
||||
"task-configs",
|
||||
"%s/config.json" % task.task_id,
|
||||
agent_config.json(exclude_none=True),
|
||||
account_id=get_func_storage(),
|
||||
StorageType.config,
|
||||
)
|
||||
reboot = False
|
||||
count = 1
|
||||
|
Reference in New Issue
Block a user