primitive 0.2.58__tar.gz → 0.2.59__tar.gz

{primitive-0.2.58 → primitive-0.2.59}/.vscode/settings.json

@@ -1,14 +1,17 @@
 {
   "cSpell.words": [
     "ajob",
+    "cafile",
     "Checkin",
     "loguru",
     "machdep",
     "Mbit",
     "mbps",
+    "percpu",
     "pkey",
     "plutil",
     "procs",
+    "RABBITMQ",
     "sessionmaker",
     "testsuites",
     "untar",

{primitive-0.2.58 → primitive-0.2.59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.58
+Version: 0.2.59
 Project-URL: Documentation, https://github.com//primitivecorp/primitive-cli#readme
 Project-URL: Issues, https://github.com//primitivecorp/primitive-cli/issues
 Project-URL: Source, https://github.com//primitivecorp/primitive-cli
@@ -22,6 +22,7 @@ Requires-Dist: click
 Requires-Dist: gql[all]
 Requires-Dist: loguru
 Requires-Dist: paramiko[invoke]
+Requires-Dist: pika>=1.3.2
 Requires-Dist: psutil>=7.0.0
 Requires-Dist: rich>=13.9.4
 Requires-Dist: speedtest-cli

{primitive-0.2.58 → primitive-0.2.59}/pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
   "speedtest-cli",
   "rich>=13.9.4",
   "psutil>=7.0.0",
+  "pika>=1.3.2",
 ]
 
 [tool.uv]

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2025-present Dylan Stein <dylan@primitive.tech>
 #
 # SPDX-License-Identifier: MIT
-__version__ = "0.2.58"
+__version__ = "0.2.59"

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/client.py

@@ -5,6 +5,8 @@ from loguru import logger
 from rich.logging import RichHandler
 from rich.traceback import install
 
+from primitive.messaging.provider import MessagingProvider
+
 from .agent.actions import Agent
 from .auth.actions import Auth
 from .daemons.actions import Daemons
@@ -80,6 +82,8 @@ class Primitive:
         else:
             self.host_config = {"username": "", "token": token, "transport": transport}
 
+        self.messaging: MessagingProvider = MessagingProvider(self)
+
         self.auth: Auth = Auth(self)
         self.organizations: Organizations = Organizations(self)
         self.projects: Projects = Projects(self)

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/graphql/sdk.py

@@ -26,7 +26,7 @@ def create_session(
     if fingerprint:
         headers["x-primitive-fingerprint"] = fingerprint
 
-    transport = AIOHTTPTransport(url=url, headers=headers)
+    transport = AIOHTTPTransport(url=url, headers=headers, ssl=True)
     session = Client(
         transport=transport,
         fetch_schema_from_transport=fetch_schema_from_transport,

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/hardware/actions.py

@@ -4,15 +4,18 @@ import json
 import platform
 import subprocess
 import typing
+from dataclasses import dataclass
 from shutil import which
 from typing import Dict, List, Optional
 
 import click
+import psutil
 from aiohttp import client_exceptions
 from gql import gql
 from loguru import logger
 
 from primitive.graphql.relay import from_base64
+from primitive.messaging.provider import MESSAGE_TYPES
 from primitive.utils.memory_size import MemorySize
 
 from ..utils.auth import guard
@@ -24,6 +27,7 @@ from .graphql.mutations import (
     register_child_hardware_mutation,
     register_hardware_mutation,
     unregister_hardware_mutation,
+    hardware_certificate_create_mutation,
 )
 from .graphql.queries import (
     hardware_details,
@@ -31,6 +35,7 @@ from .graphql.queries import (
     nested_children_hardware_list,
 )
 
+
 if typing.TYPE_CHECKING:
     pass
 
@@ -40,6 +45,12 @@ from primitive.utils.actions import BaseAction
 from primitive.utils.shell import does_executable_exist
 
 
+@dataclass
+class CertificateCreateResult:
+    certificate_id: str
+    certificate_pem: str
+
+
 class Hardware(BaseAction):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
@@ -280,6 +291,48 @@ class Hardware(BaseAction):
         system_info["gpu_config"] = self._get_gpu_config()
         return system_info
 
+    @guard
+    def certificate_create(
+        self, hardware_id: str, csr_pem: str
+    ) -> CertificateCreateResult:
+        mutation = gql(hardware_certificate_create_mutation)
+        variables = {
+            "input": {
+                "hardwareId": hardware_id,
+                "csrPem": csr_pem,
+            }
+        }
+
+        if not self.primitive.session:
+            raise Exception("No active session available for certificate creation")
+
+        result = self.primitive.session.execute(
+            mutation,
+            variable_values=variables,
+            get_execution_result=True,
+        )
+
+        if result.errors:
+            message = " ".join([error.message for error in result.errors])
+            raise Exception(message)
+
+        if not result.data:
+            raise Exception(
+                "No data received from hardware certificate creation request"
+            )
+
+        hardware_certificate_create = result.data["hardwareCertificateCreate"]
+
+        if hardware_certificate_create["__typename"] == "OperationInfo":
+            messages = hardware_certificate_create["messages"]
+            message = " ".join([error["message"] for error in messages])
+            raise Exception(message)
+
+        return CertificateCreateResult(
+            certificate_id=hardware_certificate_create["id"],
+            certificate_pem=hardware_certificate_create["certificatePem"],
+        )
+
     @guard
     def register(self, organization_id: Optional[str] = None):
         system_info = self.get_system_info()
@@ -375,6 +428,24 @@ class Hardware(BaseAction):
 
         return result
 
+    def check_in(
+        self,
+        is_healthy: bool = True,
+        is_quarantined: bool = False,
+        is_available: bool = False,
+        is_online: bool = True,
+        stopping_agent: Optional[bool] = False,
+    ):
+        message = {
+            "is_healthy": is_healthy,
+            "is_quarantined": is_quarantined,
+            "is_available": is_available,
+            "is_online": is_online,
+        }
+        self.primitive.messaging.send_message(
+            message_type=MESSAGE_TYPES.CHECK_IN, message=message
+        )
+
     @guard
     def check_in_http(
         self,
@@ -667,3 +738,28 @@ class Hardware(BaseAction):
 
         except Exception as exception:
             logger.exception(f"Error checking in children: {exception}")
+
+    def push_metrics(self):
+        self.primitive.messaging.send_message(
+            message_type=MESSAGE_TYPES.METRICS,
+            message=self.get_metrics(),
+        )
+
+    def get_metrics(self):
+        cpu_percent = psutil.cpu_percent(interval=1, percpu=False)
+        virtual_memory = psutil.virtual_memory()
+        disk_usage = psutil.disk_usage("/")
+
+        metrics = {
+            "cpu_percent": cpu_percent,
+            "memory_percent": virtual_memory.percent,
+            "memory_total": virtual_memory.total,
+            "memory_available": virtual_memory.available,
+            "memory_used": virtual_memory.used,
+            "memory_free": virtual_memory.free,
+            "disk_percent": disk_usage.percent,
+            "disk_total": disk_usage.total,
+            "disk_used": disk_usage.used,
+            "disk_free": disk_usage.free,
+        }
+        return metrics

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/hardware/commands.py

@@ -1,15 +1,18 @@
+from time import sleep
 from typing import TYPE_CHECKING, Optional
 
 import click
-
-from ..utils.printer import print_result
-from .ui import render_hardware_table
+from loguru import logger
+from primitive.utils.printer import print_result
+from primitive.hardware.ui import render_hardware_table
+from primitive.utils.x509 import (
+    generate_csr_pem,
+    write_certificate_pem,
+)
 
 if TYPE_CHECKING:
     from ..client import Primitive
 
-from loguru import logger
-
 
 @click.group()
 @click.pass_context
@@ -33,8 +36,17 @@ def systeminfo_command(context):
     type=str,
     help="Organization [slug] to register hardware with",
 )
+@click.option(
+    "--issue-certificate",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help="Issue certificate.",
+)
 @click.pass_context
-def register_command(context, organization: Optional[str] = None):
+def register_command(
+    context, organization: Optional[str] = None, issue_certificate: bool = False
+):
     """Register Hardware with Primitive"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
 
@@ -47,14 +59,30 @@ def register_command(context, organization: Optional[str] = None):
         logger.info("Registering hardware with the default organization.")
 
     result = primitive.hardware.register(organization_id=organization_id)
-    color = "green" if result else "red"
-    if result.data.get("registerHardware"):
-        message = "Hardware registered successfully"
-    else:
-        message = (
-            "There was an error registering this device. Please review the above logs."
+    hardware = result.data.get("registerHardware")
+
+    if not hardware:
+        print_result(
+            fg="red",
+            context=context,
+            message="There was an error registering this device. Please review the above logs.",
         )
-    print_result(message=message, context=context, fg=color)
+        return
+
+    if issue_certificate:
+        certificate = primitive.hardware.certificate_create(
+            hardware_id=hardware["id"],
+            csr_pem=generate_csr_pem(
+                hardware_id=hardware["pk"],
+            ),
+        )
+        write_certificate_pem(certificate.certificate_pem)
+
+    print_result(
+        fg="green",
+        context=context,
+        message="Hardware registered successfully.",
+    )
 
 
 @cli.command("unregister")
@@ -77,12 +105,7 @@ def unregister_command(context):
 def checkin_command(context):
     """Checkin Hardware with Primitive"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    check_in_http_result = primitive.hardware.check_in_http()
-    if messages := check_in_http_result.data.get("checkIn").get("messages"):
-        print_result(message=messages, context=context, fg="yellow")
-    else:
-        message = "Hardware checked in successfully"
-        print_result(message=message, context=context, fg="green")
+    primitive.hardware.check_in()
 
 
 @cli.command("list")
@@ -125,3 +148,17 @@ def get_command(context, hardware_identifier: str) -> None:
         return
     else:
         render_hardware_table([hardware])
+
+
+@cli.command("metrics")
+@click.pass_context
+@click.option("--watch", is_flag=True, help="Watch hardware metrics")
+def metrics_command(context, watch: bool = False) -> None:
+    """Get Hardware Metrics"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    if watch:
+        while True:
+            print_result(message=primitive.hardware.get_metrics(), context=context)
+            sleep(1)
+    else:
+        print_result(message=primitive.hardware.get_metrics(), context=context)

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/hardware/graphql/mutations.py

@@ -1,11 +1,29 @@
 from primitive.graphql.utility_fragments import operation_info_fragment
 
+hardware_certificate_create_mutation = (
+    operation_info_fragment
+    + """
+mutation hardwareCertificateCreate($input: HardwareCertificateCreateInput!) {
+    hardwareCertificateCreate(input: $input) {
+        __typename
+        ... on HardwareCertificate {
+            id
+            certificatePem
+        }
+        ...OperationInfoFragment
+    }
+}
+"""
+)
+
 register_hardware_mutation = (
     operation_info_fragment
     + """
 mutation registerHardware($input: RegisterHardwareInput!) {
     registerHardware(input: $input) {
         ... on Hardware {
+            id
+            pk
             fingerprint
         }
         ...OperationInfoFragment

primitive-0.2.59/src/primitive/messaging/provider.py

@@ -0,0 +1,127 @@
+import enum
+import json
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+from uuid import uuid4
+import pika
+from pika import credentials
+from ..utils.x509 import (
+    are_certificate_files_present,
+    create_ssl_context,
+    read_certificate_common_name,
+)
+from loguru import logger
+
+from primitive.utils.actions import BaseAction
+
+if TYPE_CHECKING:
+    import primitive.client
+
+EXCHANGE = "hardware"
+ROUTING_KEY = "hardware"
+VIRTUAL_HOST = "primitive"
+CELERY_TASK_NAME = "hardware.tasks.task_receive_hardware_message"
+
+
+class MESSAGE_TYPES(enum.Enum):
+    CHECK_IN = "CHECK_IN"
+    METRICS = "METRICS"
+
+
+class MessagingProvider(BaseAction):
+    def __init__(self, primitive: "primitive.client.Primitive") -> None:
+        super().__init__(primitive=primitive)
+        self.ready = False
+
+        self.fingerprint = self.primitive.host_config.get("fingerprint", None)
+        if not self.fingerprint:
+            return
+        self.token = self.primitive.host_config.get("token", None)
+        if not self.token:
+            return
+
+        rabbitmq_host = "rabbitmq-cluster.primitive.tech"
+        RABBITMQ_PORT = 443
+
+        if primitive.host == "api.dev.primitive.tech":
+            rabbitmq_host = "rabbitmq-cluster.dev.primitive.tech"
+        elif primitive.host == "api.primitive.tech":
+            rabbitmq_host = "rabbitmq-cluster.primitive.tech"
+        elif primitive.host == "api.staging.primitive.tech":
+            rabbitmq_host = "rabbitmq-cluster.staging.primitive.tech"
+        elif primitive.host == "api.test.primitive.tech":
+            rabbitmq_host = "rabbitmq-cluster.test.primitive.tech"
+        elif primitive.host == "localhost:8000":
+            rabbitmq_host = "localhost"
+            RABBITMQ_PORT = 5671
+
+        if not are_certificate_files_present():
+            logger.warning(
+                "Certificate files not present or incomplete. MessagingProvider not initialized."
+            )
+            return
+
+        ssl_context = create_ssl_context()
+        ssl_options = pika.SSLOptions(ssl_context)
+        self.common_name = read_certificate_common_name()
+
+        self.parameters = pika.ConnectionParameters(
+            host=rabbitmq_host,
+            port=RABBITMQ_PORT,
+            virtual_host=VIRTUAL_HOST,
+            ssl_options=ssl_options,
+            credentials=credentials.ExternalCredentials(),
+        )
+
+        self.ready = True
+
+    def send_message(self, message_type: MESSAGE_TYPES, message: dict[str, any]):  # type: ignore
+        if not self.ready:
+            logger.warning(
+                "send_message: cannot send message. MessagingProvider not initialized."
+            )
+            return
+
+        body = {
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "message_type": message_type.value,
+            "message": message,
+        }
+
+        full_body_for_celery = [
+            [],
+            body,
+            {"callbacks": None, "errbacks": None, "chain": None, "chord": None},
+        ]
+        with pika.BlockingConnection(parameters=self.parameters) as conn:
+            message_uuid = str(uuid4())
+            channel = conn.channel()
+
+            headers = {
+                "fingerprint": self.fingerprint,
+                "token": self.token,
+                "argsrepr": "()",
+                "id": message_uuid,
+                "ignore_result": False,
+                "kwargsrepr": str(body),
+                "lang": "py",
+                "replaced_task_nesting": 0,
+                "retries": 0,
+                "root_id": message_uuid,
+                "task": CELERY_TASK_NAME,
+            }
+
+            channel.basic_publish(
+                exchange=EXCHANGE,
+                routing_key=ROUTING_KEY,
+                body=json.dumps(full_body_for_celery),
+                properties=pika.BasicProperties(
+                    user_id=self.common_name,
+                    correlation_id=message_uuid,
+                    priority=0,
+                    delivery_mode=2,
+                    headers=headers,
+                    content_encoding="utf-8",
+                    content_type="application/json",
+                ),
+            )

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/monitor/actions.py

@@ -67,6 +67,8 @@ class Monitor(BaseAction):
                 # handles cleanup of old reservations
                 # obtains an active JobRun's ID
                 if not RUNNING_IN_CONTAINER:
+                    self.primitive.hardware.push_metrics()
+
                     hardware = self.primitive.hardware.get_own_hardware_details()
                     # fetch the latest hardware and activeReservation details
                     if active_reservation_data := hardware["activeReservation"]:

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/utils/actions.py

@@ -1,6 +1,6 @@
-import typing
+from typing import TYPE_CHECKING
 
-if typing.TYPE_CHECKING:
+if TYPE_CHECKING:
     import primitive.client
 
 

{primitive-0.2.58 → primitive-0.2.59}/src/primitive/utils/printer.py

@@ -4,7 +4,7 @@ import click
 from loguru import logger
 
 
-def print_result(message: str, context: click.Context, fg: str = None):
+def print_result(message: str | dict, context: click.Context, fg: str = None):
     """Print message to stdout or stderr"""
     if context.obj["DEBUG"]:
         logger.info(json.dumps(message))

primitive-0.2.59/src/primitive/utils/x509.py

@@ -0,0 +1,140 @@
+import ssl
+from pathlib import Path
+from loguru import logger
+
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography import x509
+from cryptography.x509.oid import NameOID
+
+PrivateKey = ec.EllipticCurvePrivateKey
+
+HOME_DIRECTORY = Path.home()
+
+PRIVATE_KEY_PATH = Path(HOME_DIRECTORY / ".config" / "primitive" / "private-key.pem")
+CERTIFICATE_PATH = Path(HOME_DIRECTORY / ".config" / "primitive" / "certificate.pem")
+
+# NOTE Only used with self-signed server certificate
+# THIS IS FOR LOCAL TESTING
+SELF_SIGNED_SERVER_CA_PATH = Path(
+    HOME_DIRECTORY / ".config" / "primitive" / "server-ca.crt.pem"
+)
+
+
+def are_certificate_files_present() -> bool:
+    return PRIVATE_KEY_PATH.exists() and CERTIFICATE_PATH.exists()
+
+
+def generate_private_key() -> PrivateKey:
+    return ec.generate_private_key(
+        curve=ec.SECP521R1(),
+    )
+
+
+def read_certificate() -> x509.Certificate:
+    cert = x509.load_pem_x509_certificate(CERTIFICATE_PATH.read_bytes())
+    return cert
+
+
+def read_certificate_common_name() -> str:
+    cert = read_certificate()
+
+    names = cert.subject.get_attributes_for_oid(NameOID.COMMON_NAME)
+
+    return str(names[0].value)
+
+
+def create_ssl_context() -> ssl.SSLContext:
+    context = ssl.create_default_context(
+        cafile=SELF_SIGNED_SERVER_CA_PATH
+        if SELF_SIGNED_SERVER_CA_PATH.exists()
+        else None,
+        purpose=ssl.Purpose.SERVER_AUTH,
+    )
+    context.load_cert_chain(
+        certfile=CERTIFICATE_PATH,
+        keyfile=PRIVATE_KEY_PATH,
+    )
+
+    return context
+
+
+def read_private_key() -> PrivateKey:
+    private_key = serialization.load_pem_private_key(
+        data=PRIVATE_KEY_PATH.read_bytes(),
+        password=None,
+    )
+
+    if not isinstance(private_key, PrivateKey):
+        raise Exception(
+            f"Expected private key type {PrivateKey.__name__}, got {type(private_key).__name__}"
+        )
+
+    return private_key
+
+
+def ensure_private_key() -> PrivateKey:
+    private_key_path = PRIVATE_KEY_PATH
+
+    if private_key_path.exists():
+        private_key = read_private_key()
+    else:
+        logger.info("Generating private key.")
+
+        private_key = generate_private_key()
+
+        with private_key_path.open("wb") as f:
+            f.write(
+                private_key.private_bytes(
+                    encoding=serialization.Encoding.PEM,
+                    format=serialization.PrivateFormat.PKCS8,
+                    encryption_algorithm=serialization.NoEncryption(),
+                )
+            )
+
+    return private_key
+
+
+def write_certificate_pem(certificate_pem: str):
+    CERTIFICATE_PATH.write_text(
+        data=certificate_pem,
+        encoding="utf-8",
+    )
+
+
+def check_certificate():
+    try:
+        if CERTIFICATE_PATH.exists():
+            with CERTIFICATE_PATH.open("rb") as file:
+                private_key = read_private_key()
+                crt = x509.load_pem_x509_certificate(file.read())
+
+                # NOTE: Make sure certificate match private key
+                return crt.public_key() == private_key.public_key()
+    except Exception:
+        return False
+
+    return False
+
+
+def generate_csr_pem(hardware_id: str) -> str:
+    builder = x509.CertificateSigningRequestBuilder()
+    builder = builder.subject_name(
+        x509.Name(
+            [
+                x509.NameAttribute(NameOID.COMMON_NAME, hardware_id),
+            ]
+        )
+    )
+    builder = builder.add_extension(
+        x509.BasicConstraints(ca=False, path_length=None),
+        critical=True,
+    )
+    csr = builder.sign(
+        algorithm=hashes.SHA512(),
+        private_key=ensure_private_key(),
+    )
+
+    return csr.public_bytes(encoding=serialization.Encoding.PEM).decode(
+        encoding="utf-8", errors="strict"
+    )

{primitive-0.2.58 → primitive-0.2.59}/uv.lock

@@ -801,6 +801,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/9e/c3/059298687310d527a58bb01f3b1965787ee3b40dce76752eda8b44e9a2c5/pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523", size = 63772, upload-time = "2023-11-25T06:56:14.81Z" },
 ]
 
+[[package]]
+name = "pika"
+version = "1.3.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/db/db/d4102f356af18f316c67f2cead8ece307f731dd63140e2c71f170ddacf9b/pika-1.3.2.tar.gz", hash = "sha256:b2a327ddddf8570b4965b3576ac77091b850262d34ce8c1d8cb4e4146aa4145f", size = 145029, upload-time = "2023-05-05T14:25:43.368Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f9/f3/f412836ec714d36f0f4ab581b84c491e3f42c6b5b97a6c6ed1817f3c16d0/pika-1.3.2-py3-none-any.whl", hash = "sha256:0779a7c1fafd805672796085560d290213a465e4f6f76a6fb19e378d8041a14f", size = 155415, upload-time = "2023-05-05T14:25:41.484Z" },
+]
+
 [[package]]
 name = "primitive"
 source = { editable = "." }
@@ -809,6 +818,7 @@ dependencies = [
     { name = "gql", extra = ["all"] },
     { name = "loguru" },
     { name = "paramiko", extra = ["invoke"] },
+    { name = "pika" },
     { name = "psutil" },
     { name = "rich" },
     { name = "speedtest-cli" },
@@ -828,6 +838,7 @@ requires-dist = [
     { name = "gql", extras = ["all"] },
     { name = "loguru" },
     { name = "paramiko", extras = ["invoke"] },
+    { name = "pika", specifier = ">=1.3.2" },
     { name = "psutil", specifier = ">=7.0.0" },
     { name = "rich", specifier = ">=13.9.4" },
     { name = "speedtest-cli" },