util.py

#!/bin/env python3
import os
import subprocess
import time
import socket

from params import *

from typing import List, Dict


def run_remote(server, directory, cmd, user=REMOTE_USER, bg=False):
    cmd = f"ssh -o StrictHostKeyChecking=no {user}@{server} 'cd {directory} && {cmd}'"

    if bg:
        try:
            proc = 0
            proc = subprocess.Popen(
                cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
            )
        except Exception as e:
            print(f"An error occurred: {e}")
            return -1
        return proc
    else:
        exit_code = os.system(cmd)
        if exit_code != 0:
            print(f"Command {cmd} failed with exit code: {exit_code}")
        return exit_code


def copy_to_remote(server, src, dst, user=REMOTE_USER, rec=False):
    cmd = "scp "
    if rec:
        cmd += "-r "
    cmd += f"{src} {user}@{server}:{dst}"
    os.system(cmd)


def copy_from_remote(server, src, dst, user=REMOTE_USER, rec=False):
    cmd = "scp "
    if rec:
        cmd += "-r "
    cmd = f"scp {user}@{server}:{src} {dst}"
    os.system(cmd)


def make(server):
    run_remote(server, f"~/{REPO_NAME}", "make")
    run_remote(server, f"~/{REPO_NAME}/apps/sird_tester", "make")


def make_initial(server):
    run_remote(
        server,
        "~",
        f"cd {REPO_NAME} && make submodules && make -j && cd ksched && make",
    )

def get_caladan_ip(server):
    if "sm110p-10s106" in server:
        i = int(server.split("sm110p-10s106")[1])
        return f"192.168.1.{i}"
    elif "amd" in server:
        i = int(server.split("amd")[1])
        return f"192.168.1.{i}"
    else:
        raise ValueError(f"No manual caladan ip for {server}")


IP_MAP_CACHE = {} # {server}-{interface} -> address

def get_interface_ip(server, interface):
    ''' Return interface IP for both Caladan and SSH interfaces.'''
    if interface == get_caladan_iface(server):
        try:
            return get_caladan_ip(server)
        except:
            pass
    print(
        f"Ip for {server} and interface {interface} not found in static map. Attempting to get it by sshing into {server}"
    )
    key = f"{server}-{interface}"
    if key in IP_MAP_CACHE:
        return IP_MAP_CACHE[key]
    cmd = f"ip addr show {interface} | grep inet"
    cmd = f'ssh -o StrictHostKeyChecking=no {server} "{cmd}"'
    result = subprocess.check_output(cmd, shell=True, text=True).strip()
    result = result.split("inet")[1].split("metric")[0].strip()
    result = result.split("/")[0].strip()
    IP_MAP_CACHE[key] = result
    return result


def get_other_ips(machine_ip, all_ips):
    other_servers = all_ips.copy()
    other_servers.remove(machine_ip)
    other_servers = " ".join(other_servers)
    return other_servers

def is_ipv4(ip):
    """Check if the given string is a valid IPv4 address."""
    try:
        socket.inet_pton(socket.AF_INET, ip)
        return True
    except socket.error:
        return False

def is_ipv6(ip):
    """Check if the given string is a valid IPv6 address."""
    try:
        socket.inet_pton(socket.AF_INET6, ip)
        return True
    except socket.error:
        return False

def is_ip_address(ip):
    """Check if the given string is a valid IPv4 or IPv6 address."""
    return is_ipv4(ip) or is_ipv6(ip)