import urllib.request, json
from pymongo import MongoClient

# Flip to strict
db = MongoClient('mongodb://localhost:27017')['rmm_db']
db.config.update_one({'_id': 'agent_auth'}, {'$set': {'mode': 'strict'}}, upsert=True)

def test_heartbeat(token=None):
    headers = {'Content-Type': 'application/json'}
    if token is not None:
        headers['X-Agent-Token'] = token
    req = urllib.request.Request(
        'http://localhost:8000/api/heartbeat?client_id=TEST_AUTH',
        data=json.dumps({'cpu_percent':10,'cpu_temp':40,'memory_percent':50,'memory_total_gb':16,'memory_free_gb':8,'disks':[],'gpus':[],'agent_version':'test'}).encode('utf-8'),
        headers=headers
    )
    try:
        urllib.request.urlopen(req)
        return 'OK'
    except Exception as e:
        return str(e)

print('Strict mode:')
print('No token:', test_heartbeat())
print('Wrong token:', test_heartbeat('wrong'))
print('Correct token:', test_heartbeat('170LRiZafLZThrIGQ6VlP8K9FPw7i8InwzM3dVoYnK4'))

# Dispatch update_agent to JAGAN
db.clients.update_one({'_id': 'JAGAN'}, {'$set': {'pending_command': 'update_agent'}})
print('Dispatched update_agent to JAGAN!')

# Flip back to grace so we don't break other things during the migration
db.config.update_one({'_id': 'agent_auth'}, {'$set': {'mode': 'grace'}})