idlers-agent/main.py

import os
import time
import json
import requests

import psutil
import tldextract
import dns.resolver

# get ram and disk in MB and GB respectively
def get_ram_and_disk():
    ram = psutil.virtual_memory().total >> 20
    disk = psutil.disk_usage('/').total >> 30
    return ram, disk

def get_cpu_count():
    return psutil.cpu_count()

def get_bandwidth():
    return 2000

def get_hostname():
    return os.uname().nodename

def get_public_ip():
    try:
        response = requests.get('https://api.ipify.org')
        response.raise_for_status()
        return response.text
    except:
        return '127.0.0.1'

def get_domain():
    extracted = tldextract.extract(get_hostname())
    return "{}.{}".format(extracted.domain, extracted.suffix)

def get_nameservers():
    try:
        answers = dns.resolver.resolve(get_domain(), 'NS')
        return [str(rdata) for rdata in answers]
    except Exception as e:
        return []

def get_os():
    return 27

# Create POST data
def create_post_data():
    ram, disk = get_ram_and_disk()
    nameservers = get_nameservers()
    post_data = {
        "server_type": 1,
        "os_id": get_os(),
        "provider_id": 10,
        "location_id": 15,
        "ssh_port": 22,
        "ram": ram >> 10,   # convert to GB
        "ram_as_mb": ram,   # in MB
        "disk": disk,       # in GB
        "disk_as_gb": disk, # in GB
        "cpu": get_cpu_count(),
        "bandwidth": get_bandwidth(),
        "was_promo": 1,
        "active": 1,
        "show_public": 0,
        "owned_since": "2022-01-01",
        "ram_type": "GB",
        "disk_type": "GB",
        "currency": "USD",
        "price": 4,
        "payment_term": 1,
        "hostname": get_hostname(),
        "next_due_date": "2022-02-01",

        # Above ones are required fields, below are optional
        "ns1": nameservers[0] if nameservers else "",
        "ns2": nameservers[1] if len(nameservers) > 1 else "",
        # "has_yabs": 0,
        "ip1": get_public_ip(), # empty string or None gives 400 Bad Request
        # "as_usd": 4,
        # "usd_per_month": 4,
        }
    return post_data

def get_existing_servers(host, api_key):
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ' + api_key,
    }

    try:
        response = requests.get(host, headers=headers)
        response.raise_for_status()
        return response.json()
    except:
        return []

def existing_server_id(existing_servers, post_data):
    for server in existing_servers:
        if server['hostname'] == post_data['hostname']:
            return server['id']
    return None

def create_server(host, api_key, post_data):
    # Create the headers for the request
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ' + api_key,
    }
    payload = json.dumps(post_data)

    # Send the POST request
    try:
        response = requests.post(host, headers=headers, data=payload)
        response.raise_for_status()
    except requests.exceptions.HTTPError as errh:
        print ("Http Error:",errh)
    except requests.exceptions.ConnectionError as errc:
        print ("Error Connecting:",errc)
    except requests.exceptions.Timeout as errt:
        print ("Timeout Error:",errt)
    except requests.exceptions.RequestException as err:
        print ("Something went wrong:",err)
    else:
        print(response.text)

def update_server(host, api_key, post_data, server_id):
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ' + api_key,
    }

    # remove following keys from post_data
    # ssh_port, ip1, currency, price, payment_term, next_due_date
    post_data.pop('ssh_port', None)
    post_data.pop('ip1', None)
    post_data.pop('currency', None)
    post_data.pop('price', None)
    post_data.pop('payment_term', None)
    post_data.pop('next_due_date', None)

    payload = json.dumps(post_data)

    try:
        response = requests.put(host + '/' + str(server_id), headers=headers, data=payload)
        response.raise_for_status()
    except requests.exceptions.HTTPError as errh:
        print ("Http Error:",errh)
    except requests.exceptions.ConnectionError as errc:
        print ("Error Connecting:",errc)
    except requests.exceptions.Timeout as errt:
        print ("Timeout Error:",errt)
    except requests.exceptions.RequestException as err:
        print ("Something went wrong:",err)
    else:
        print(response.text)

def send_post_request():
    # Get the API key and HOST from environment variables
    api_key = os.getenv('API_KEY')
    host = os.getenv('HOST')
    if not api_key:
        raise Exception('API_KEY not found in environment variables')
    if not host:
        raise Exception('HOST not found in environment variables')

    # Create the data for the post request
    post_data = create_post_data()

    # Get the existing servers
    existing_servers = get_existing_servers(host, api_key)

    # Check if the server already exists
    server_id = existing_server_id(existing_servers, post_data)

    # If the server exists, update it
    if server_id:
        print('Server already exists with id: {}, Updating...'.format(server_id))
        update_server(host, api_key, post_data, server_id)
    else:
        print('Server does not exist, Creating...')
        create_server(host, api_key, post_data)

if __name__ == '__main__':
    send_post_request()