idlers-agent/agent.py

import os
import urllib.request
import logging
import json
import http.client

NON_UPDATABLE_KEYS = [
    'server_type',
    'os_id',
    'provider_id',
    'location_id',
    'ssh_port', 
    'bandwidth', 
    'was_promo',
    'active',
    'owned_since', 
    'currency', 
    'price', 
    'payment_term', 
    'next_due_date',
    'ip1'
]

class ServerData:
    def __init__(self):
        self.hostname = os.uname().nodename
        self.public_ip = self.get_public_ip()
        logging.basicConfig(level=logging.INFO)

    def get_ram_and_disk(self):
        with open('/proc/meminfo', 'r') as f:
            meminfo = f.read()
        ram = int([x for x in meminfo.split('\n') if 'MemTotal' in x][0].split()[1]) // 1024
        with open('/proc/diskstats', 'r') as f:
            diskstats = f.read()
        disk = sum(int(x.split()[9]) for x in diskstats.split('\n') if x) * 512 // 10**9
        logging.info(f"RAM: {ram}MB, Disk: {disk}GB")
        return ram, disk

    def get_cpu_count(self):
        cpu_count = 0
        if os.path.isfile('/usr/sbin/dmidecode'):
            try:
                output = subprocess.check_output(['sudo', '/usr/sbin/dmidecode', '-t', 'processor']).decode('utf-8')
                core_match = re.search(r'Core Count: (\d+)', output)
                thread_match = re.search(r'Thread Count: (\d+)', output)
                if core_match and thread_match:
                    cpu_count = int(core_match.group(1)) * int(thread_match.group(1))
            except subprocess.CalledProcessError:
                pass
        if cpu_count == 0:
            with open('/proc/cpuinfo', 'r') as f:
                cpuinfo = f.read()
            cpu_count = cpuinfo.count('processor')
        logging.info(f"CPU Count: {cpu_count}")
        return cpu_count

    def get_bandwidth(self):
        bandwidth = 2000
        logging.info(f"Bandwidth: {bandwidth}")
        return bandwidth

    def get_public_ip(self):
        try:
            response = urllib.request.urlopen('https://api.ipify.org')
            return response.read().decode()
        except Exception as e:
            logging.error(f"Failed to get public IP: {e}")
            return '127.0.0.1'

    def get_os(self):
        os_id = 27
        logging.info(f"OS ID: {os_id}")
        return os_id

    def create_post_data(self):
        ram, disk = self.get_ram_and_disk()
        post_data = {
            "server_type": 1,
            "os_id": self.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": self.get_cpu_count(),
            "bandwidth": self.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": self.hostname,
            "next_due_date": "2022-02-01",
            "ip1": self.public_ip,
        }
        logging.info("Post data created")
        return post_data


class ServerManager:
    def __init__(self, host, api_key):
        self.host = host
        self.api_key = api_key
        self.headers = {
            'Content-Type': 'application/json',
            'Authorization': 'Bearer ' + self.api_key,
        }
        logging.basicConfig(level=logging.INFO)

    def send_request(self, method, endpoint, data=None):
        url = self.host + endpoint
        payload = json.dumps(data).encode('utf-8') if data else None
        req = urllib.request.Request(url, data=payload, headers=self.headers, method=method)

        try:
            with urllib.request.urlopen(req) as response:
                if method == 'GET':
                    return json.loads(response.read().decode())
                else:
                    return response.read().decode()
        except urllib.error.HTTPError as e:
            logging.error(f"Request failed with {e}")
            raise

    def get_existing_servers(self):
        return self.send_request('GET', '/api/servers')

    def create_server(self, post_data):
        logging.info("Creating server...")
        return self.send_request('POST', '/api/servers', post_data)

    def update_server(self, post_data, server_id):
        # remove following keys from post_data
        for key in NON_UPDATABLE_KEYS:
            post_data.pop(key, None)
        logging.info(f"Updating server with id {server_id}...")
        return self.send_request('PUT', '/api/servers/' + str(server_id), post_data)

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

def validate_env_vars():
    api_key = os.getenv('AGENT_API')
    host = os.getenv('HOST')
    if not api_key:
        raise Exception('AGENT_API not found in environment variables')
    if not host:
        raise Exception('HOST not found in environment variables')
    return host, api_key

def main():
    logging.basicConfig(level=logging.INFO)

    host, api_key = validate_env_vars()

    server_data = ServerData()
    post_data = server_data.create_post_data()

    server_manager = ServerManager(host, api_key)

    # Check if the server already exists
    server_id = server_manager.existing_server_id(post_data)

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

if __name__ == '__main__':
    main()