Untitled

mail@pastecode.io avatar
unknown
python
2 years ago
8.8 kB
4
Indexable
Never
import os
import json
import paramiko
from scp import SCPClient
from pathlib import *


def to_windows(file_name: str):
    return file_name.replace("/", "\\")


def to_linux(file_name: str):
    linux_name = file_name.replace("\\", "/")
    linux_name = linux_name.replace("/\./", "/")
    return linux_name


# =============================================================================
class SSH(object):

    hostname = ''
    username = ''
    password = ''
    remote_root = ''
    remote_dir = ''

    # =============================================================================
    def __init__(self, root: str):
        defaults_file = os.environ['AT_DEFAULTS']
        json_file = os.path.join(os.path.dirname(defaults_file), "cluster.json")
        with open(json_file) as f:
            auth = json.load(f)
        self.hostname = auth["hostname"]
        self.username = auth["username"]
        self.password = auth["password"]
        self.remote_root = root
        self.remote_dir = self.remote_root
        self.ssh = paramiko.SSHClient()
        self.ssh.load_system_host_keys()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        # self.ssh.connect(self.hostname, username=self.username, password=self.password, port=22, timeout=3600000000)
        self.ssh.connect(self.hostname, username=self.username, password=self.password, port=22)
        self.execute("mkdir -p " + self.remote_root)
        ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("mkdir -p " + self.remote_dir)

    # =============================================================================
    def close(self):
        self.ssh.close()

    # =============================================================================
    def create_remote_dir(self, new_dir: str):
        ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("mkdir -p " + new_dir)
        print("Create remote dir " + new_dir)
        if ssh_stderr:
            print(ssh_stderr)
        else:
            print("Ok")

    # =============================================================================
    def change_remote_dir(self, dir_name):
        self.remote_dir = self.remote_root + "/" + dir_name  # assuming remote is a Linux system

    # =============================================================================
    def upload(self, file_list, target_dir):
        print("************** Upload ***********************")
        with SCPClient(self.ssh.get_transport()) as scp:
            for f in file_list:
                source = to_windows(f.strip())
                print("Uploading ", source, "from", os.getcwd(), " to ", target_dir)
                scp.put(source, target_dir)

    # =============================================================================
    def download(self, file_list):
        print("************** Download *********************")
        # ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("pwd")
        with SCPClient(self.ssh.get_transport()) as scp:
            for f in file_list:
                source = self.remote_dir + "/" + f
                target = os.getcwd()
                print("Downloading ", source, " to ", target)
                scp.get(source, target)

    # =============================================================================
    def execute(self, remote_command: str = "", silent: bool = False):
        full_command = 'cd ' + self.remote_dir + "; " + remote_command
        ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command("pwd")
        pwd = ssh_stdout.read().decode().strip()
        ssh_std_in, ssh_stdout, ssh_stderr = self.ssh.exec_command(full_command)
        output = ssh_stdout.read().decode()
        error = ssh_stderr.read().decode()
        if not silent:
            print("************** Execute ***********************")
            print(self.username, "@", self.hostname, ":", pwd)
            print(full_command)
            print("************** Output  ***********************")
            print(output)
            if error:
                print("************** Error  ************************")
                print(error)
            # print("************** Done   ***********************")


# =============================================================================
class RemoteExecution(object):
    LOCAL_BASE_DIR = None
    LOCAL_SOURCE_DIR = None

    REMOTE_ROOT_DIR = None
    REMOTE_BASE_DIR = None
    REMOTE_SOURCE_DIR = None

    # =========================================================================
    def __init__(self, inp_data):
        self.LOCAL_BASE_DIR = inp_data["Environment"]["BASE_DIR"]
        self.LOCAL_SOURCE_DIR = os.environ["AT_ROOT"]

        self.LOCAL_BASE_DIR = self.LOCAL_BASE_DIR.replace("\\", "/")
        split = r'proc/'
        self.REMOTE_ROOT_DIR = inp_data["RemoteServer"]["root_dir"]
        try:
            self.REMOTE_BASE_DIR = os.path.join(self.REMOTE_ROOT_DIR, self.LOCAL_BASE_DIR.split(split)[1])
        except IndexError:
            self.REMOTE_BASE_DIR = os.path.join(self.REMOTE_ROOT_DIR, os.path.basename(self.LOCAL_BASE_DIR))
        self.REMOTE_BASE_DIR = self.REMOTE_BASE_DIR.replace("\\", "/")
        print("Remote Base Dir: ", self.REMOTE_BASE_DIR)
        self.REMOTE_SOURCE_DIR = inp_data["RemoteServer"]["source_dir"]
        self.source_file_filters = inp_data["RemoteServer"]["source_file_filter"]

    # =============================================================================
    def init_remote(self):
        print("Initializing remote session")
        uploads = [str(x) for x in Path(".").iterdir() if x.is_file()]
        print("DBG hhh: Remote Base Dir: ", self.REMOTE_BASE_DIR)
        ssh = SSH(self.REMOTE_BASE_DIR)
        ssh.upload(uploads, self.REMOTE_BASE_DIR)
        ssh.close()

    # =============================================================================
    def upload_sources(self):
        print("Uploading sources")
        os.chdir(self.LOCAL_SOURCE_DIR)
        ssh = SSH(self.REMOTE_SOURCE_DIR)
        for file_filter in self.source_file_filters:
            for root, dirs, files in os.walk("."):
                for file in files:
                    if file.endswith(".py"):
                        if "*" in file_filter or file_filter in file:
                            source = os.path.join(root, file)
                            target = os.path.join(self.REMOTE_SOURCE_DIR, source)
                            target_dir = str(Path(target).parent).replace("\\", "/")
                            ssh.upload([source], target_dir)
        ssh.close()
        os.chdir(self.LOCAL_BASE_DIR)

    # =============================================================================
    def run_remote(self, step_name, step_dir, uploads, downloads, upload_all=False):
        print("Running step_name " + step_name + " remotely")

        REMOTE_STEP_DIR = os.path.join(self.REMOTE_BASE_DIR, step_dir).replace("\\", "/")

        ssh = SSH(self.REMOTE_BASE_DIR)
        ssh.create_remote_dir(REMOTE_STEP_DIR)
        # ssh.upload(["../"+self.REPORT_FILE], self.REMOTE_BASE_DIR)
        if upload_all:
            uploads = [str(x) for x in Path(".").iterdir() if x.is_file()]
        ssh.upload(uploads, REMOTE_STEP_DIR)
        if step_name == "ALL":
            ssh.execute("python $AT_ROOT/runCase.py --nogui   >&  runCase.log")
        else:
            ssh.execute("python $AT_ROOT/runCase.py --nogui --slave --input_file runCase.inp.yml --step=" + step_name + " >& " + step_name + ".log")
        ssh.change_remote_dir(step_dir)
        ssh.download(downloads)
        ssh.close()

    # =============================================================================
    def download_file(self, file):
        ssh = SSH(self.REMOTE_BASE_DIR)
        ssh.download([file])


# ====================================================
# Template for Main
# ====================================================
if __name__ == "__main__":
    from Lib.user_arguments import UserOptions
    usr = UserOptions()
    usr.input_file.EXTENSION = "*.inp"
    usr.input_file.DESCRIPTION = "Input file [runCase.inp] "
    usr.input_file.REQUIRED = True
    usr.input_file.DEFAULT = "runCase.inp"
    usr.work_file.REQUIRED = False
    usr.work_file.ENABLE_GUI = False
    usr.get_options()
    inp_file = usr.input_file.NAME
    os.chdir(usr.input_file.DIRECTORY)

    rem = SSH()
    rem.change_remote_dir("Transfer_speed")
    rem.upload([usr.input_file.BASENAME])
    rem.execute("cp geom1.inp mima.inp")
    rem.execute("ls")
    rem.download(["mima.inp"])