arthexis 0.1.3__py3-none-any.whl

nodes/admin.py

@@ -0,0 +1,347 @@
+from django.contrib import admin, messages
+from django.urls import path, reverse
+from django.shortcuts import redirect, render
+from django.utils.html import format_html
+from django import forms
+from django.contrib.admin.widgets import FilteredSelectMultiple
+from app.widgets import CopyColorWidget, CodeEditorWidget
+from django.db import models
+from django.conf import settings
+from pathlib import Path
+from django.http import HttpResponse
+import base64
+import pyperclip
+from pyperclip import PyperclipException
+import uuid
+from .utils import capture_screenshot, save_screenshot
+from .actions import NodeAction
+
+from .models import (
+    Node,
+    EmailOutbox,
+    NodeRole,
+    ContentSample,
+    NodeTask,
+    NetMessage,
+    User,
+)
+from core.admin import UserAdmin as CoreUserAdmin
+
+
+class NodeAdminForm(forms.ModelForm):
+    class Meta:
+        model = Node
+        fields = "__all__"
+        widgets = {"badge_color": CopyColorWidget()}
+
+
+@admin.register(Node)
+class NodeAdmin(admin.ModelAdmin):
+    list_display = (
+        "hostname",
+        "mac_address",
+        "address",
+        "port",
+        "role",
+        "last_seen",
+    )
+    search_fields = ("hostname", "address", "mac_address")
+    change_list_template = "admin/nodes/node/change_list.html"
+    change_form_template = "admin/nodes/node/change_form.html"
+    form = NodeAdminForm
+    actions = ["run_task", "take_screenshots"]
+
+
+    def get_urls(self):
+        urls = super().get_urls()
+        custom = [
+            path(
+                "register-current/",
+                self.admin_site.admin_view(self.register_current),
+                name="nodes_node_register_current",
+            ),
+            path(
+                "<int:node_id>/action/<str:action>/",
+                self.admin_site.admin_view(self.action_view),
+                name="nodes_node_action",
+            ),
+            path(
+                "<int:node_id>/public-key/",
+                self.admin_site.admin_view(self.public_key),
+                name="nodes_node_public_key",
+            ),
+        ]
+        return custom + urls
+
+    def register_current(self, request):
+        """Create or update this host and offer browser node registration."""
+        node, created = Node.register_current()
+        if created:
+            self.message_user(
+                request, f"Current host registered as {node}", messages.SUCCESS
+            )
+        token = uuid.uuid4().hex
+        context = {
+            "token": token,
+            "register_url": reverse("register-node"),
+        }
+        return render(request, "admin/nodes/node/register_remote.html", context)
+
+    def public_key(self, request, node_id):
+        node = self.get_object(request, node_id)
+        if not node:
+            self.message_user(request, "Unknown node", messages.ERROR)
+            return redirect("..")
+        security_dir = Path(settings.BASE_DIR) / "security"
+        pub_path = security_dir / f"{node.public_endpoint}.pub"
+        if pub_path.exists():
+            response = HttpResponse(pub_path.read_bytes(), content_type="text/plain")
+            response["Content-Disposition"] = f'attachment; filename="{pub_path.name}"'
+            return response
+        self.message_user(request, "Public key not found", messages.ERROR)
+        return redirect("..")
+
+    def run_task(self, request, queryset):
+        if "apply" in request.POST:
+            recipe_text = request.POST.get("recipe", "")
+            task_obj, _ = NodeTask.objects.get_or_create(recipe=recipe_text)
+            results = []
+            for node in queryset:
+                try:
+                    output = task_obj.run(node)
+                except Exception as exc:
+                    output = str(exc)
+                results.append((node, output))
+            context = {"recipe": recipe_text, "results": results}
+            return render(request, "admin/nodes/task_result.html", context)
+        context = {"nodes": queryset}
+        return render(request, "admin/nodes/node/run_task.html", context)
+
+    run_task.short_description = "Run task"
+
+    @admin.action(description="Take Screenshots")
+    def take_screenshots(self, request, queryset):
+        tx = uuid.uuid4()
+        sources = getattr(settings, "SCREENSHOT_SOURCES", ["/"])
+        count = 0
+        for node in queryset:
+            for source in sources:
+                try:
+                    url = source.format(node=node, address=node.address, port=node.port)
+                except Exception:
+                    url = source
+                if not url.startswith("http"):
+                    url = f"http://{node.address}:{node.port}{url}"
+                try:
+                    path = capture_screenshot(url)
+                except Exception as exc:  # pragma: no cover - selenium issues
+                    self.message_user(request, f"{node}: {exc}", messages.ERROR)
+                    continue
+                sample = save_screenshot(
+                    path, node=node, method="ADMIN", transaction_uuid=tx
+                )
+                if sample:
+                    count += 1
+        self.message_user(request, f"{count} screenshots captured", messages.SUCCESS)
+
+    def changeform_view(self, request, object_id=None, form_url="", extra_context=None):
+        extra_context = extra_context or {}
+        extra_context["node_actions"] = NodeAction.get_actions()
+        if object_id:
+            extra_context["public_key_url"] = reverse(
+                "admin:nodes_node_public_key", args=[object_id]
+            )
+        return super().changeform_view(
+            request, object_id, form_url, extra_context=extra_context
+        )
+
+    def action_view(self, request, node_id, action):
+        node = self.get_object(request, node_id)
+        action_cls = NodeAction.registry.get(action)
+        if not node or not action_cls:
+            self.message_user(request, "Unknown node action", messages.ERROR)
+            return redirect("..")
+        try:
+            result = action_cls.run(node)
+            if hasattr(result, "status_code"):
+                return result
+            self.message_user(
+                request,
+                f"{action_cls.display_name} executed successfully",
+                messages.SUCCESS,
+            )
+        except NotImplementedError:
+            self.message_user(
+                request,
+                "Remote node actions are not yet implemented",
+                messages.WARNING,
+            )
+        except Exception as exc:  # pragma: no cover - unexpected errors
+            self.message_user(request, str(exc), messages.ERROR)
+        return redirect(reverse("admin:nodes_node_change", args=[node_id]))
+
+
+@admin.register(EmailOutbox)
+class EmailOutboxAdmin(admin.ModelAdmin):
+    list_display = ("node", "host", "port", "username", "use_tls", "use_ssl")
+
+
+class NodeRoleAdminForm(forms.ModelForm):
+    nodes = forms.ModelMultipleChoiceField(
+        queryset=Node.objects.all(),
+        required=False,
+        widget=FilteredSelectMultiple("Nodes", False),
+    )
+
+    class Meta:
+        model = NodeRole
+        fields = ("name", "description", "nodes")
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.instance.pk:
+            self.fields["nodes"].initial = self.instance.node_set.all()
+
+
+@admin.register(NodeRole)
+class NodeRoleAdmin(admin.ModelAdmin):
+    form = NodeRoleAdminForm
+    list_display = ("name", "description")
+
+    def save_model(self, request, obj, form, change):
+        obj.node_set.set(form.cleaned_data.get("nodes", []))
+
+
+@admin.register(ContentSample)
+class ContentSampleAdmin(admin.ModelAdmin):
+    list_display = ("name", "kind", "node", "user", "created_at")
+    readonly_fields = ("created_at", "name", "user", "image_preview")
+
+    def get_urls(self):
+        urls = super().get_urls()
+        custom = [
+            path(
+                "from-clipboard/",
+                self.admin_site.admin_view(self.add_from_clipboard),
+                name="nodes_contentsample_from_clipboard",
+            ),
+            path(
+                "capture/",
+                self.admin_site.admin_view(self.capture_now),
+                name="nodes_contentsample_capture",
+            ),
+        ]
+        return custom + urls
+
+    def add_from_clipboard(self, request):
+        try:
+            content = pyperclip.paste()
+        except PyperclipException as exc:  # pragma: no cover - depends on OS clipboard
+            self.message_user(request, f"Clipboard error: {exc}", level=messages.ERROR)
+            return redirect("..")
+        if not content:
+            self.message_user(request, "Clipboard is empty.", level=messages.INFO)
+            return redirect("..")
+        if ContentSample.objects.filter(content=content, kind=ContentSample.TEXT).exists():
+            self.message_user(
+                request, "Duplicate sample not created.", level=messages.INFO
+            )
+            return redirect("..")
+        user = request.user if request.user.is_authenticated else None
+        ContentSample.objects.create(content=content, user=user, kind=ContentSample.TEXT)
+        self.message_user(
+            request, "Text sample added from clipboard.", level=messages.SUCCESS
+        )
+        return redirect("..")
+
+    def capture_now(self, request):
+        node = Node.get_local()
+        url = request.build_absolute_uri("/")
+        try:
+            path = capture_screenshot(url)
+        except Exception as exc:  # pragma: no cover - depends on selenium setup
+            self.message_user(request, str(exc), level=messages.ERROR)
+            return redirect("..")
+        sample = save_screenshot(path, node=node, method="ADMIN")
+        if sample:
+            self.message_user(request, f"Screenshot saved to {path}", messages.SUCCESS)
+        else:
+            self.message_user(
+                request, "Duplicate screenshot; not saved", messages.INFO
+            )
+        return redirect("..")
+
+    @admin.display(description="Screenshot")
+    def image_preview(self, obj):
+        if not obj or obj.kind != ContentSample.IMAGE or not obj.path:
+            return ""
+        file_path = Path(obj.path)
+        if not file_path.is_absolute():
+            file_path = settings.LOG_DIR / file_path
+        if not file_path.exists():
+            return "File not found"
+        with file_path.open("rb") as f:
+            encoded = base64.b64encode(f.read()).decode("ascii")
+        return format_html(
+            '<img src="data:image/png;base64,{}" style="max-width:100%;" />',
+            encoded,
+        )
+
+
+@admin.register(NetMessage)
+class NetMessageAdmin(admin.ModelAdmin):
+    list_display = ("subject", "body", "reach", "created", "complete")
+    search_fields = ("subject", "body")
+    list_filter = ("complete", "reach")
+    ordering = ("-created",)
+    readonly_fields = ("complete",)
+    actions = ["send_messages"]
+
+    def send_messages(self, request, queryset):
+        for msg in queryset:
+            msg.propagate()
+        self.message_user(request, f"{queryset.count()} messages sent")
+
+    send_messages.short_description = "Send selected messages"
+
+
+class NodeTaskForm(forms.ModelForm):
+    class Meta:
+        model = NodeTask
+        fields = "__all__"
+        widgets = {"recipe": CodeEditorWidget()}
+
+
+@admin.register(NodeTask)
+class NodeTaskAdmin(admin.ModelAdmin):
+    form = NodeTaskForm
+    list_display = ("recipe", "role", "created")
+    actions = ["execute"]
+
+    def execute(self, request, queryset):
+        if queryset.count() != 1:
+            self.message_user(
+                request, "Please select exactly one task", messages.ERROR
+            )
+            return
+        task_obj = queryset.first()
+        if "apply" in request.POST:
+            node_ids = request.POST.getlist("nodes")
+            nodes_qs = Node.objects.filter(pk__in=node_ids)
+            results = []
+            for node in nodes_qs:
+                try:
+                    output = task_obj.run(node)
+                except Exception as exc:
+                    output = str(exc)
+                results.append((node, output))
+            context = {"recipe": task_obj.recipe, "results": results}
+            return render(request, "admin/nodes/task_result.html", context)
+        nodes = Node.objects.all()
+        context = {"nodes": nodes, "task_obj": task_obj}
+        return render(request, "admin/nodes/nodetask/run.html", context)
+
+    execute.short_description = "Run task on nodes"
+
+
+admin.site.register(User, CoreUserAdmin)

nodes/apps.py

@@ -0,0 +1,76 @@
+import logging
+import os
+import socket
+import threading
+import time
+from pathlib import Path
+
+from django.apps import AppConfig
+from django.conf import settings
+from django.core.signals import request_started
+from django.db import connections
+from django.db.utils import OperationalError
+from utils import revision
+
+
+logger = logging.getLogger(__name__)
+
+
+def _startup_notification() -> None:
+    """Queue a notification with host:port and version on a background thread."""
+
+    host = socket.gethostname()
+    try:
+        address = socket.gethostbyname(host)
+    except socket.gaierror:
+        address = host
+
+    port = os.environ.get("PORT", "8000")
+
+    version = ""
+    ver_path = Path(settings.BASE_DIR) / "VERSION"
+    if ver_path.exists():
+        version = ver_path.read_text().strip()
+
+    revision_value = revision.get_revision()
+    rev_short = revision_value[-6:] if revision_value else ""
+
+    body = f"v{version}"
+    if rev_short:
+        body += f" r{rev_short}"
+
+    def _worker() -> None:  # pragma: no cover - background thread
+        # Allow the LCD a moment to become ready and retry a few times
+        for _ in range(5):
+            try:
+                from nodes.models import NetMessage
+
+                NetMessage.broadcast(subject=f"{address}:{port}", body=body)
+                break
+            except Exception:
+                time.sleep(1)
+
+    threading.Thread(target=_worker, name="startup-notify", daemon=True).start()
+
+
+def _trigger_startup_notification(**_: object) -> None:
+    """Send the startup notification once a request has started."""
+
+    request_started.disconnect(_trigger_startup_notification, dispatch_uid="nodes-startup")
+    try:
+        connections["default"].ensure_connection()
+    except OperationalError:
+        logger.exception("Startup notification skipped: database unavailable")
+        return
+    _startup_notification()
+
+
+class NodesConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "nodes"
+    verbose_name = "Node Infrastructure"
+
+    def ready(self):  # pragma: no cover - exercised on app start
+        request_started.connect(
+            _trigger_startup_notification, dispatch_uid="nodes-startup"
+        )

nodes/lcd.py

@@ -0,0 +1,151 @@
+"""Minimal driver for PCF8574/PCF8574A I2C LCD1602 displays.
+
+The implementation is adapted from the example provided in the
+instructions.  It is intentionally lightweight and only implements the
+operations required for this project: initialisation, clearing the
+screen and writing text to a specific position.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import time
+from dataclasses import dataclass
+from typing import List
+
+try:  # pragma: no cover - hardware dependent
+    import smbus  # type: ignore
+except Exception:  # pragma: no cover - missing dependency
+    try:  # pragma: no cover - hardware dependent
+        import smbus2 as smbus  # type: ignore
+    except Exception:  # pragma: no cover - missing dependency
+        smbus = None  # type: ignore
+
+
+class LCDUnavailableError(RuntimeError):
+    """Raised when the LCD cannot be initialised."""
+
+
+@dataclass
+class _BusWrapper:
+    """Wrapper around :class:`smbus.SMBus` to allow mocking in tests."""
+
+    channel: int
+
+    def write_byte(self, addr: int, data: int) -> None:  # pragma: no cover - thin wrapper
+        if smbus is None:
+            raise LCDUnavailableError("smbus not available")
+        bus = smbus.SMBus(self.channel)
+        bus.write_byte(addr, data)
+        bus.close()
+
+
+class CharLCD1602:
+    """Minimal driver for PCF8574/PCF8574A I2C backpack (LCD1602)."""
+
+    def __init__(self, bus: _BusWrapper | None = None) -> None:
+        if smbus is None:  # pragma: no cover - hardware dependent
+            raise LCDUnavailableError("smbus not available")
+        self.bus = bus or _BusWrapper(1)
+        self.BLEN = 1
+        self.PCF8574_address = 0x27
+        self.PCF8574A_address = 0x3F
+        self.LCD_ADDR = self.PCF8574_address
+
+    def _write_word(self, addr: int, data: int) -> None:
+        if self.BLEN:
+            data |= 0x08
+        else:
+            data &= 0xF7
+        self.bus.write_byte(addr, data)
+
+    def _pulse_enable(self, data: int) -> None:
+        self._write_word(self.LCD_ADDR, data | 0x04)
+        time.sleep(0.0005)
+        self._write_word(self.LCD_ADDR, data & ~0x04)
+        time.sleep(0.0001)
+
+    def send_command(self, cmd: int) -> None:
+        high = cmd & 0xF0
+        low = (cmd << 4) & 0xF0
+        self._write_word(self.LCD_ADDR, high)
+        self._pulse_enable(high)
+        self._write_word(self.LCD_ADDR, low)
+        self._pulse_enable(low)
+        # Give the LCD time to process the command to avoid garbled output.
+        time.sleep(0.001)
+
+    def send_data(self, data: int) -> None:
+        high = (data & 0xF0) | 0x01
+        low = ((data << 4) & 0xF0) | 0x01
+        self._write_word(self.LCD_ADDR, high)
+        self._pulse_enable(high)
+        self._write_word(self.LCD_ADDR, low)
+        self._pulse_enable(low)
+        # Allow the LCD controller to catch up between data writes.
+        time.sleep(0.001)
+
+    def i2c_scan(self) -> List[str]:  # pragma: no cover - requires hardware
+        """Return a list of detected I2C addresses.
+
+        The implementation relies on the external ``i2cdetect`` command.  On
+        systems where ``i2c-tools`` is not installed or the command cannot be
+        executed (e.g. insufficient permissions), the function returns an empty
+        list so callers can fall back to a sensible default address.
+        """
+
+        cmd = "i2cdetect -y 1 | awk 'NR>1 {$1=\"\"; print}'"
+        try:
+            out = subprocess.check_output(cmd, shell=True).decode()
+        except Exception:  # pragma: no cover - depends on environment
+            return []
+        out = out.replace("\n", "").replace(" --", "")
+        return [tok for tok in out.split(" ") if tok]
+
+    def init_lcd(self, addr: int | None = None, bl: int = 1) -> None:
+        self.BLEN = 1 if bl else 0
+        if addr is None:
+            try:
+                found = self.i2c_scan()
+            except Exception:  # pragma: no cover - i2c detection issues
+                found = []
+            if "3f" in found or "3F" in found:
+                self.LCD_ADDR = self.PCF8574A_address
+            else:
+                # Default to the common PCF8574 address (0x27) when detection
+                # fails or returns no recognised addresses. This mirrors the
+                # behaviour prior to introducing automatic address detection and
+                # prevents the display from remaining uninitialised on systems
+                # without ``i2c-tools``.
+                self.LCD_ADDR = self.PCF8574_address
+        else:
+            self.LCD_ADDR = addr
+
+        time.sleep(0.05)
+        self.send_command(0x33)
+        self.send_command(0x32)
+        self.send_command(0x28)
+        self.send_command(0x0C)
+        self.send_command(0x06)
+        self.clear()
+        self._write_word(self.LCD_ADDR, 0x00)
+
+    def clear(self) -> None:
+        self.send_command(0x01)
+        time.sleep(0.002)
+
+    def reset(self) -> None:
+        """Re-run the initialisation sequence to recover the display."""
+        self.init_lcd(addr=self.LCD_ADDR, bl=self.BLEN)
+
+    def set_backlight(self, on: bool = True) -> None:  # pragma: no cover - hardware dependent
+        self.BLEN = 1 if on else 0
+        self._write_word(self.LCD_ADDR, 0x00)
+
+    def write(self, x: int, y: int, s: str) -> None:
+        x = max(0, min(15, int(x)))
+        y = 0 if int(y) <= 0 else 1
+        addr = 0x80 + 0x40 * y + x
+        self.send_command(addr)
+        for ch in str(s):
+            self.send_data(ord(ch))