idlers-agent/agent.py

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

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()
        self.dmidecode_data = self.parse_dmidecode_output()
        logging.basicConfig(level=logging.INFO)

    def parse_dmidecode_output(self):
        '''
        Example dmidecode output:

        Handle 0x0069, DMI type 20, 35 bytes
        Memory Device Mapped Address
            Starting Address: 0x00600000000
            Ending Address: 0x007FFFFFFFF
            Range Size: 8 GB
            Physical Device Handle: 0x0067
            Memory Array Mapped Address Handle: 0x006A
            Partition Row Position: Unknown
            Interleave Position: 2
            Interleaved Data Depth: 2

        Handle 0x006A, DMI type 19, 31 bytes
        Memory Array Mapped Address
            Starting Address: 0x00000000000
            Ending Address: 0x007FFFFFFFF
            Range Size: 32 GB
            Physical Array Handle: 0x005E
            Partition Width: 4
        '''
        if os.path.isfile('/usr/sbin/dmidecode'):
            try:
                # ignore error messages produced by the command
                output = subprocess.check_output(['sudo', '/usr/sbin/dmidecode'], stderr=subprocess.DEVNULL).decode('utf-8')

                # each section is separated by two newlines
                sections = output.split('\n\n')
                parsed_sections = []
                for section in sections:
                    
                    # each line in a section is separated by a newline
                    lines = section.split('\n')

                    # first line contains the DMI type - only need number
                    match = re.search(r'DMI type (\d+)', lines[0])
                    if match:
                        dmi_type = int(match.group(1))
                    else:   # skip if no DMI type in this section
                        continue

                    # each section is a dictionary with DMIType as key
                    section_dict = {'DMIType': dmi_type}
                    if len(lines) > 1:
                        section_dict['description'] = lines[1].strip()

                    for line in lines[2:]:  # skip first two lines - already processed

                        # each line has tabs at the beginning and space(maybe optional) between key and value
                        match = re.match(r'\t(.+):\s*(.*)', line)
                        if match:
                            key, value = match.groups()
                            section_dict[key] = value
                    parsed_sections.append(section_dict)
                return parsed_sections
            except subprocess.CalledProcessError:
                pass
        return []

    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
        for section in self.dmidecode_data:
            if section['DMIType'] == 4:  # 4 corresponds to processor
                core_count = int(section.get('Core Count', '0'))
                thread_count = int(section.get('Thread Count', '0'))
                cpu_count = core_count * thread_count
        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

    def create_note_data(self):
        chassis_info = None
        for section in self.dmidecode_data:
            if section['DMIType'] == 1:
                chassis_info = section
                break
        if chassis_info:
            chassis_model = chassis_info.get('Product Name', 'Unknown')
            chassis_serial = chassis_info.get('Serial Number', 'Unknown')
        else:
            chassis_model = chassis_serial = 'Unknown'

        processor_info = [section for section in self.dmidecode_data if section['DMIType'] == 4]
        processor_model = processor_info[0].get('Version', 'Unknown') if processor_info else 'Unknown'
        processor_count = len(processor_info)

        ram_info = [section for section in self.dmidecode_data if section['DMIType'] == 17]
        ram_details = []
        for ram in ram_info:
            size = ram.get('Size', 'Unknown')
            speed = ram.get('Speed', 'Unknown')
            configured_speed = ram.get('Configured Memory Speed', 'Unknown')
            total_width = int(ram.get('Total Width', "0").split()[0])
            data_width = int(ram.get('Data Width', "0").split()[0])
            ecc = 'Yes' if total_width > data_width else 'No'
            serial_number = ram.get('Serial Number', 'Unknown')
            ram_type = ram.get('Type', 'Unknown')
            ram_details.append("Size: {}, Speed: {} @ {}, ECC: {}, Serial Number: {}, Type: {}".format(size, speed, configured_speed, ecc, serial_number, ram_type))

        note = "Chassis Model: {} | Serial Number: {} ||| Processor Model: {} | Count: {} ||| RAM Details: {}".format(
        chassis_model, chassis_serial, processor_model, processor_count, ' | '.join(ram_details))

        note_data = {
            'note': note,
        }
        return note_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 get_note(self, service_id):
        return self.send_request('GET', '/api/notes/' + str(service_id))

    def create_note(self, post_data):
        logging.info("Creating note...")
        return self.send_request('POST', '/api/notes', post_data)

    def update_note(self, post_data, service_id):
        logging.info(f"Updating note with id {service_id}...")
        return self.send_request('PUT', '/api/notes/' + str(service_id), post_data)

    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 upsert_server(self, post_data):
        server_id = self.existing_server_id(post_data)

        if server_id:
            logging.info('Server already exists with id: {}, Updating...'.format(server_id))
            response = self.update_server(post_data, server_id)
        else:
            logging.info('Server does not exist, Creating...')
            response = self.create_server(post_data)

        # Extract the server_id from the response
        server_id = json.loads(response).get('server_id', None)
        if server_id is None:
            logging.error('Failed to get server_id from response: {}'.format(response))
            raise ValueError('Failed to get server_id from response')

        return server_id
    
    def upsert_note(self, note_data, server_id):
        note_data['service_id'] = server_id
        try:
            note = self.get_note(server_id)
        except urllib.error.HTTPError:
            note = None

        if note:
            return self.update_note(note_data, server_id)
        else:
            return self.create_note(note_data)

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)

    server_id = server_manager.upsert_server(post_data)
    logging.info('Server id: {}'.format(server_id))

    note_data = server_data.create_note_data()
    server_manager.upsert_note(note_data, server_id)

if __name__ == '__main__':
    main()