shopdb-flask/scripts/dev/seed_endpoints_sample.py

"""Seed sample endpoint telemetry for the Fleet module (DEV ONLY).

Attaches EndpointState snapshots plus a handful of curated EndpointEvents to
existing Computer assets so the fleet list, the PC-detail asset panels
(status / services / recent events), and the dashboard quiet + stability cards
have something to render before a real agent exists.

Sample rows are marked with a sentinel agentversion ("0.0.0-sample") so they can
be listed and cleared without touching real agent data.

DEV ONLY: this file lives under scripts/dev/ and must never ship in a
production image (see devnotes/live-container-and-testing-workflow.md). Do not
import it from application code.

Usage (in the API container):
    docker exec shopdb-flask-wilm-api-1 python scripts/dev/seed_endpoints_sample.py
    docker exec shopdb-flask-wilm-api-1 python scripts/dev/seed_endpoints_sample.py --clear
    docker exec shopdb-flask-wilm-api-1 python scripts/dev/seed_endpoints_sample.py --limit 8
"""

import argparse
import json
import os
import sys
from datetime import datetime, timedelta, timezone

# Repo root on sys.path so `shopdb` and `plugins` import when run as a script.
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))

from shopdb import create_app  # noqa: E402
from shopdb.api import db  # noqa: E402

SAMPLE_MARKER = '0.0.0-sample'


def _naive_utcnow():
    return datetime.now(timezone.utc).replace(tzinfo=None)


def _services(*running_then_stopped):
    """Build a services snapshot. Positional pairs are (name, state)."""
    return [{'name': name, 'state': state} for name, state in running_then_stopped]


def _ports(*pairs):
    return [{'port': port, 'process': process} for port, process in pairs]


def _interfaces(ip, mac):
    return [{'name': 'Ethernet', 'ip': ip, 'mac': mac}]


# Per-machine telemetry shapes. minutesago drives whether a machine reads as
# fresh, stale, or quiet; a couple are pushed past a plausible quiet window so
# the dashboard quiet card has rows. Everything is offline-plausible because
# there is no live agent.
SAMPLE_PROFILES = [
    {
        'loggedinuser': 'CORP\\jmartin',
        'agentversion': SAMPLE_MARKER,
        'minutesago': 3,
        'uptimeseconds': 4 * 86400 + 6 * 3600,
        'idleseconds': 12 * 60,
        'cpupercent': 14.0,
        'memorypercent': 46.0,
        'services': _services(('Spooler', 'Running'), ('W32Time', 'Running'),
                              ('WinRM', 'Running'), ('MSSQLSERVER', 'Running')),
        'ports': _ports((3389, 'svchost'), (5985, 'winrm'), (445, 'System')),
        'interfaces': _interfaces('10.20.4.11', '00:1A:2B:3C:4D:01'),
        'events': [(4624, -8), (4672, -8), (4698, -120)],
    },
    {
        'loggedinuser': 'CORP\\operator2',
        'agentversion': SAMPLE_MARKER,
        'minutesago': 9,
        'uptimeseconds': 2 * 3600 + 40 * 60,
        'idleseconds': 3 * 60,
        'cpupercent': 61.0,
        'memorypercent': 72.0,
        'services': _services(('Spooler', 'Running'), ('W32Time', 'Running'),
                              ('WinRM', 'Stopped')),
        'ports': _ports((3389, 'svchost'), (445, 'System')),
        'interfaces': _interfaces('10.20.4.27', '00:1A:2B:3C:4D:02'),
        # Recent unexpected shutdown -> shows on the stability card.
        'events': [(41, -50), (1001, -30), (4625, -12)],
    },
    {
        'loggedinuser': None,
        'agentversion': SAMPLE_MARKER,
        'minutesago': 22,
        'uptimeseconds': 11 * 86400,
        'idleseconds': 9 * 3600,
        'cpupercent': 3.0,
        'memorypercent': 38.0,
        'services': _services(('Spooler', 'Running'), ('W32Time', 'Running'),
                              ('WinRM', 'Running')),
        'ports': _ports((3389, 'svchost'), (5985, 'winrm')),
        'interfaces': _interfaces('10.20.4.44', '00:1A:2B:3C:4D:03'),
        'events': [(4624, -600), (7011, -180)],
    },
    {
        'loggedinuser': 'CORP\\nightshift',
        'agentversion': SAMPLE_MARKER,
        # Quiet: no report in over a day.
        'minutesago': 1500,
        'uptimeseconds': 30 * 60,
        'idleseconds': 30 * 60,
        'cpupercent': 1.0,
        'memorypercent': 29.0,
        'services': _services(('Spooler', 'Stopped'), ('W32Time', 'Running')),
        'ports': _ports((445, 'System'),),
        'interfaces': _interfaces('10.20.4.58', '00:1A:2B:3C:4D:04'),
        'events': [(55, -1600), (41, -1580)],
    },
    {
        'loggedinuser': 'CORP\\qa_kiosk',
        'agentversion': SAMPLE_MARKER,
        'minutesago': 6,
        'uptimeseconds': 6 * 3600,
        'idleseconds': 45 * 60,
        'cpupercent': 22.0,
        'memorypercent': 55.0,
        'services': _services(('Spooler', 'Running'), ('W32Time', 'Running'),
                              ('WinRM', 'Running'), ('BITS', 'Running')),
        'ports': _ports((3389, 'svchost'), (5985, 'winrm'), (135, 'rpc')),
        'interfaces': _interfaces('10.20.4.63', '00:1A:2B:3C:4D:05'),
        'events': [(4624, -20), (10016, -90), (4697, -240)],
    },
]


def _clear(EndpointState, EndpointEvent):
    sample_states = EndpointState.query.filter_by(agentversion=SAMPLE_MARKER).all()
    asset_ids = [s.assetid for s in sample_states]
    events_removed = 0
    if asset_ids:
        events_removed = (EndpointEvent.query
                          .filter(EndpointEvent.assetid.in_(asset_ids))
                          .delete(synchronize_session=False))
    states_removed = (EndpointState.query
                      .filter_by(agentversion=SAMPLE_MARKER)
                      .delete(synchronize_session=False))
    db.session.commit()
    print('Cleared {} sample state(s) and {} sample event(s).'
          .format(states_removed, events_removed))


def _seed(limit):
    from plugins.computers.models import Computer
    from plugins.endpoints.models import EndpointState, EndpointEvent
    from plugins.endpoints.eventcatalog import EVENT_CATALOG

    computers = (Computer.query
                 .filter(Computer.assetid.isnot(None))
                 .order_by(Computer.computerid.asc())
                 .limit(limit)
                 .all())
    if not computers:
        print('No Computer assets found to attach sample endpoints to. '
              'Seed the computers plugin first.')
        return

    now = _naive_utcnow()
    created, updated, events_added = 0, 0, 0

    for index, computer in enumerate(computers):
        profile = SAMPLE_PROFILES[index % len(SAMPLE_PROFILES)]
        reported = now - timedelta(minutes=profile['minutesago'])
        boot = reported - timedelta(seconds=profile['uptimeseconds'])

        state = EndpointState.query.filter_by(assetid=computer.assetid).first()
        if state and state.agentversion not in (None, SAMPLE_MARKER):
            # Never stomp a real agent's snapshot.
            continue
        if not state:
            state = EndpointState(assetid=computer.assetid, createdat=now)
            db.session.add(state)
            created += 1
        else:
            updated += 1

        state.hostname = computer.hostname or 'HOST-{}'.format(computer.computerid)
        state.loggedinuser = profile['loggedinuser']
        state.agentversion = SAMPLE_MARKER
        state.lastboottime = boot
        state.uptimeseconds = profile['uptimeseconds']
        state.idleseconds = profile['idleseconds']
        state.cpupercent = profile['cpupercent']
        state.memorypercent = profile['memorypercent']
        state.servicesjson = json.dumps(profile['services'])
        state.portsjson = json.dumps(profile['ports'])
        state.interfacesjson = json.dumps(profile['interfaces'])
        state.scheduledtasksjson = json.dumps([])
        state.lastreportedat = reported

        # Refresh this asset's sample events so re-running stays idempotent.
        EndpointEvent.query.filter_by(assetid=computer.assetid).delete(
            synchronize_session=False)
        for eventid, offset_minutes in profile['events']:
            title, severity, channel = EVENT_CATALOG.get(
                eventid, ('Event {}'.format(eventid), 'info', 'System'))
            event = EndpointEvent(
                assetid=computer.assetid,
                eventid=eventid,
                eventchannel=channel,
                eventsource='sample-seed',
                eventtime=reported + timedelta(minutes=offset_minutes),
                eventmessage='{} (sample)'.format(title),
                payloadjson=json.dumps({'sample': True, 'severity': severity}),
                recordedat=now,
            )
            db.session.add(event)
            events_added += 1

    db.session.commit()
    print('Seeded sample endpoints: {} created, {} updated, {} event(s) added '
          'across {} computer(s).'
          .format(created, updated, events_added, len(computers)))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--clear', action='store_true',
                        help='Remove sample-marked endpoint states and their events.')
    parser.add_argument('--limit', type=int, default=5,
                        help='Max number of computers to attach sample endpoints to.')
    args = parser.parse_args()

    app = create_app(os.environ.get('FLASK_ENV', 'production'))
    with app.app_context():
        from plugins.endpoints.models import EndpointState, EndpointEvent
        if args.clear:
            _clear(EndpointState, EndpointEvent)
        else:
            _seed(args.limit)


if __name__ == '__main__':
    main()