34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
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'}})
|