mirror of
https://github.com/microsoft/onefuzz.git
synced 2025-06-17 12:28:07 +00:00
initial public release
This commit is contained in:
Brian Caswell
273
src/api-service/__app__/onefuzzlib/azure/vm.py
Normal file
273
src/api-service/__app__/onefuzzlib/azure/vm.py
Normal file
@ -0,0 +1,273 @@
|
||||
#!/usr/bin/env python
|
||||
#
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# Licensed under the MIT License.
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional, Union, cast
|
||||
from uuid import UUID
|
||||
|
||||
from azure.mgmt.compute import ComputeManagementClient
|
||||
from azure.mgmt.compute.models import VirtualMachine
|
||||
from msrestazure.azure_exceptions import CloudError
|
||||
from onefuzztypes.enums import OS, ErrorCode
|
||||
from onefuzztypes.models import Authentication, Error
|
||||
from onefuzztypes.primitives import Extension, Region
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .creds import get_base_resource_group, mgmt_client_factory
|
||||
from .disk import delete_disk, list_disks
|
||||
from .image import get_os
|
||||
from .ip import create_public_nic, delete_ip, delete_nic, get_ip, get_public_nic
|
||||
|
||||
|
||||
def get_vm(name: str) -> Optional[VirtualMachine]:
|
||||
resource_group = get_base_resource_group()
|
||||
|
||||
logging.debug("getting vm: %s %s - %s", resource_group, name)
|
||||
compute_client = mgmt_client_factory(ComputeManagementClient)
|
||||
try:
|
||||
return cast(
|
||||
VirtualMachine,
|
||||
compute_client.virtual_machines.get(
|
||||
resource_group, name, expand="instanceView"
|
||||
),
|
||||
)
|
||||
except CloudError as err:
|
||||
logging.debug("vm does not exist %s", err)
|
||||
return None
|
||||
|
||||
|
||||
def create_vm(
|
||||
name: str,
|
||||
location: str,
|
||||
vm_sku: str,
|
||||
image: str,
|
||||
password: str,
|
||||
ssh_public_key: str,
|
||||
) -> Union[None, Error]:
|
||||
resource_group = get_base_resource_group()
|
||||
logging.info("creating vm %s:%s:%s", resource_group, location, name)
|
||||
|
||||
compute_client = mgmt_client_factory(ComputeManagementClient)
|
||||
|
||||
nic = get_public_nic(resource_group, name)
|
||||
if nic is None:
|
||||
result = create_public_nic(resource_group, name, location)
|
||||
if isinstance(result, Error):
|
||||
return result
|
||||
logging.info("waiting on nic creation")
|
||||
return None
|
||||
|
||||
if image.startswith("/"):
|
||||
image_ref = {"id": image}
|
||||
else:
|
||||
image_val = image.split(":", 4)
|
||||
image_ref = {
|
||||
"publisher": image_val[0],
|
||||
"offer": image_val[1],
|
||||
"sku": image_val[2],
|
||||
"version": image_val[3],
|
||||
}
|
||||
|
||||
params: Dict = {
|
||||
"location": location,
|
||||
"os_profile": {
|
||||
"computer_name": "node",
|
||||
"admin_username": "onefuzz",
|
||||
"admin_password": password,
|
||||
},
|
||||
"hardware_profile": {"vm_size": vm_sku},
|
||||
"storage_profile": {"image_reference": image_ref},
|
||||
"network_profile": {"network_interfaces": [{"id": nic.id}]},
|
||||
}
|
||||
|
||||
image_os = get_os(location, image)
|
||||
if isinstance(image_os, Error):
|
||||
return image_os
|
||||
|
||||
if image_os == OS.linux:
|
||||
params["os_profile"]["linux_configuration"] = {
|
||||
"disable_password_authentication": True,
|
||||
"ssh": {
|
||||
"public_keys": [
|
||||
{
|
||||
"path": "/home/onefuzz/.ssh/authorized_keys",
|
||||
"key_data": ssh_public_key,
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
if "ONEFUZZ_OWNER" in os.environ:
|
||||
params["tags"] = {"OWNER": os.environ["ONEFUZZ_OWNER"]}
|
||||
|
||||
try:
|
||||
compute_client.virtual_machines.create_or_update(resource_group, name, params)
|
||||
except CloudError as err:
|
||||
if "The request failed due to conflict with a concurrent request" in str(err):
|
||||
logging.debug(
|
||||
"create VM had conflicts with concurrent request, ignoring %s", err
|
||||
)
|
||||
return None
|
||||
return Error(code=ErrorCode.VM_CREATE_FAILED, errors=[str(err)])
|
||||
return None
|
||||
|
||||
|
||||
def get_extension(vm_name: str, extension_name: str) -> Optional[Any]:
|
||||
resource_group = get_base_resource_group()
|
||||
|
||||
logging.debug(
|
||||
"getting extension: %s:%s:%s - %s",
|
||||
resource_group,
|
||||
vm_name,
|
||||
extension_name,
|
||||
)
|
||||
compute_client = mgmt_client_factory(ComputeManagementClient)
|
||||
try:
|
||||
return compute_client.virtual_machine_extensions.get(
|
||||
resource_group, vm_name, extension_name
|
||||
)
|
||||
except CloudError as err:
|
||||
logging.error("extension does not exist %s", err)
|
||||
return None
|
||||
|
||||
|
||||
def create_extension(vm_name: str, extension: Dict) -> Any:
|
||||
resource_group = get_base_resource_group()
|
||||
|
||||
logging.info(
|
||||
"creating extension: %s:%s:%s", resource_group, vm_name, extension["name"]
|
||||
)
|
||||
compute_client = mgmt_client_factory(ComputeManagementClient)
|
||||
return compute_client.virtual_machine_extensions.create_or_update(
|
||||
resource_group, vm_name, extension["name"], extension
|
||||
)
|
||||
|
||||
|
||||
def delete_vm(name: str) -> Any:
|
||||
resource_group = get_base_resource_group()
|
||||
|
||||
logging.info("deleting vm: %s %s", resource_group, name)
|
||||
compute_client = mgmt_client_factory(ComputeManagementClient)
|
||||
return compute_client.virtual_machines.delete(resource_group, name)
|
||||
|
||||
|
||||
def has_components(name: str) -> bool:
|
||||
# check if any of the components associated with a VM still exist.
|
||||
#
|
||||
# Azure VM Deletion requires we first delete the VM, then delete all of it's
|
||||
# resources. This is required to ensure we've cleaned it all up before
|
||||
# marking it "done"
|
||||
resource_group = get_base_resource_group()
|
||||
if get_vm(name):
|
||||
return True
|
||||
if get_public_nic(resource_group, name):
|
||||
return True
|
||||
if get_ip(resource_group, name):
|
||||
return True
|
||||
|
||||
disks = [x.name for x in list_disks(resource_group) if x.name.startswith(name)]
|
||||
if disks:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def delete_vm_components(name: str) -> bool:
|
||||
resource_group = get_base_resource_group()
|
||||
logging.info("deleting vm components %s:%s", resource_group, name)
|
||||
if get_vm(name):
|
||||
logging.info("deleting vm %s:%s", resource_group, name)
|
||||
delete_vm(name)
|
||||
return False
|
||||
|
||||
if get_public_nic(resource_group, name):
|
||||
logging.info("deleting nic %s:%s", resource_group, name)
|
||||
delete_nic(resource_group, name)
|
||||
return False
|
||||
|
||||
if get_ip(resource_group, name):
|
||||
logging.info("deleting ip %s:%s", resource_group, name)
|
||||
delete_ip(resource_group, name)
|
||||
return False
|
||||
|
||||
disks = [x.name for x in list_disks(resource_group) if x.name.startswith(name)]
|
||||
if disks:
|
||||
for disk in disks:
|
||||
logging.info("deleting disk %s:%s", resource_group, disk)
|
||||
delete_disk(resource_group, disk)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class VM(BaseModel):
|
||||
name: Union[UUID, str]
|
||||
region: Region
|
||||
sku: str
|
||||
image: str
|
||||
auth: Authentication
|
||||
|
||||
def is_deleted(self) -> bool:
|
||||
return has_components(str(self.name))
|
||||
|
||||
def exists(self) -> bool:
|
||||
return self.get() is not None
|
||||
|
||||
def get(self) -> Optional[VirtualMachine]:
|
||||
return get_vm(str(self.name))
|
||||
|
||||
def create(self) -> Union[None, Error]:
|
||||
if self.get() is not None:
|
||||
return None
|
||||
|
||||
logging.info("vm creating: %s", self.name)
|
||||
return create_vm(
|
||||
str(self.name),
|
||||
self.region,
|
||||
self.sku,
|
||||
self.image,
|
||||
self.auth.password,
|
||||
self.auth.public_key,
|
||||
)
|
||||
|
||||
def delete(self) -> bool:
|
||||
return delete_vm_components(str(self.name))
|
||||
|
||||
def add_extensions(self, extensions: List[Extension]) -> Union[bool, Error]:
|
||||
status = []
|
||||
to_create = []
|
||||
for config in extensions:
|
||||
if not isinstance(config["name"], str):
|
||||
logging.error("vm agent - incompatable name: %s", repr(config))
|
||||
continue
|
||||
extension = get_extension(str(self.name), config["name"])
|
||||
|
||||
if extension:
|
||||
logging.info(
|
||||
"vm extension state: %s - %s - %s",
|
||||
self.name,
|
||||
config["name"],
|
||||
extension.provisioning_state,
|
||||
)
|
||||
status.append(extension.provisioning_state)
|
||||
else:
|
||||
to_create.append(config)
|
||||
|
||||
if to_create:
|
||||
for config in to_create:
|
||||
create_extension(str(self.name), config)
|
||||
else:
|
||||
if all([x == "Succeeded" for x in status]):
|
||||
return True
|
||||
elif "Failed" in status:
|
||||
return Error(
|
||||
code=ErrorCode.VM_CREATE_FAILED,
|
||||
errors=["failed to launch extension"],
|
||||
)
|
||||
elif not ("Creating" in status or "Updating" in status):
|
||||
logging.error("vm agent - unknown state %s: %s", self.name, status)
|
||||
|
||||
return False
|
||||
Reference in New Issue
Block a user