import secrets
from django.core.management.base import BaseCommand,CommandError
from django.db import transaction
from monitoring.models import Segment,Server,AgentCredential,Audit
from monitoring.auth import digest
class Command(BaseCommand):
    help='Register host or rotate its token; prints the secret once. Keep output private.'
    def add_arguments(self,p):
        p.add_argument('name');p.add_argument('--segment',required=True);p.add_argument('--address',default='');p.add_argument('--rotate',action='store_true')
    @transaction.atomic
    def handle(self,*args,**o):
        seg,_=Segment.objects.get_or_create(name=o['segment'])
        s,created=Server.objects.get_or_create(name=o['name'],defaults={'segment':seg,'address':o['address']})
        if not created and not o['rotate']: raise CommandError('Already registered; use --rotate to revoke old credentials')
        AgentCredential.objects.filter(server=s).update(active=False)
        token=secrets.token_urlsafe(48)
        AgentCredential.objects.create(server=s,digest=digest(token))
        Audit.objects.create(actor='management-command',action='register_agent',details={'server':str(s.pk),'rotated':not created})
        self.stdout.write('SERVER_ID='+str(s.pk))
        self.stdout.write('AGENT_TOKEN='+token)