diff --git a/src/cli/onefuzz/cli.py b/src/cli/onefuzz/cli.py index ebd0c68727..7e28c5ad09 100644 --- a/src/cli/onefuzz/cli.py +++ b/src/cli/onefuzz/cli.py @@ -490,18 +490,33 @@ def print_nested_help(self, args: argparse.Namespace) -> None: level += 1 +def normalize(result: Any) -> Any: + """Convert arbitrary result streams into something filterable with jmespath""" + if isinstance(result, BaseModel): + return normalize(result.dict(exclude_none=True)) + if isinstance(result, Model): + return normalize(result.as_dict()) + if isinstance(result, list): + return [normalize(x) for x in result] + if isinstance(result, dict): + return {normalize(k): normalize(v) for (k, v) in result.items()} + if isinstance(result, Enum): + return result.name + if isinstance(result, UUID): + return str(result) + if isinstance(result, (int, float, str)): + return result + + logging.debug(f"unable to normalize type f{type(result)}") + + return result + + def output(result: Any, output_format: str, expression: Optional[Any]) -> None: if isinstance(result, bytes): sys.stdout.buffer.write(result) else: - if isinstance(result, list) and result and isinstance(result[0], BaseModel): - # cycling through json resolves all of the nested BaseModel objects - result = [json.loads(x.json(exclude_none=True)) for x in result] - if isinstance(result, BaseModel): - # cycling through json resolves all of the nested BaseModel objects - result = json.loads(result.json(exclude_none=True)) - if isinstance(result, Model): - result = result.as_dict() + result = normalize(result) if expression is not None: result = expression.search(result) if result is not None: diff --git a/src/cli/onefuzz/status/cmd.py b/src/cli/onefuzz/status/cmd.py index 7649968d59..6863062403 100644 --- a/src/cli/onefuzz/status/cmd.py +++ b/src/cli/onefuzz/status/cmd.py @@ -4,11 +4,12 @@ # Licensed under the MIT License. from collections import defaultdict -from typing import DefaultDict, List, Optional, Set +from typing import DefaultDict, Dict, List, Optional, Set from uuid import UUID -from onefuzztypes.enums import ContainerType, JobState +from onefuzztypes.enums import ContainerType, JobState, NodeState from onefuzztypes.primitives import Container +from pydantic import BaseModel from ..api import UUID_EXPANSION, Command from .cache import JobFilter @@ -20,9 +21,37 @@ def short(value: UUID) -> str: return str(value).split("-")[0] +class PoolStatus(BaseModel): + # number of nodes in each node state + node_state: Dict[NodeState, int] + # number of VMs used for each task + tasks: Dict[UUID, int] + + class Status(Command): """Monitor status of Onefuzz Instance""" + def pool(self, *, name: str) -> PoolStatus: + pool = self.onefuzz.pools.get(name) + nodes = self.onefuzz.nodes.list(pool_name=pool.name) + + tasks = {} + node_state = {} + for node in nodes: + if node.state not in node_state: + node_state[node.state] = 0 + node_state[node.state] += 1 + + if node.state not in [NodeState.free, NodeState.init]: + node = self.onefuzz.nodes.get(node.machine_id) + if node.tasks is not None: + for entry in node.tasks: + task_id, _task_state = entry + if task_id not in tasks: + tasks[task_id] = 0 + tasks[task_id] += 1 + return PoolStatus(tasks=tasks, node_state=node_state) + def project( self, *,