from paramiko import SSHClient

##---------------------------------------------------------------------------------------##
class SSH(object):
##---------------------------------------------------------------------------------------##
    ##-----------------------------------------------------------------------------------##
    def __init__(self, devices):
    ##-----------------------------------------------------------------------------------##
        for device in devices:
            self.devices[device] = {}
            self.devices[device]['user'] = devices[device]['user']
            self.devices[device]['ip'] = devices[device]['ip']

    ##-----------------------------------------------------------------------------------##
    def command(self, device, command):
    ##-----------------------------------------------------------------------------------##
        self.client = SSHClient()
        self.client.load_system_host_keys()
        self.client.connect(self.devices[device]['ip'], self.devices[device]['user'])

        if(command=="shutdown"):
            self.shutdown(device)
        if(command=="reboot"):
            self.reboot(device)
        if(command=="upgrade"):
            self.upgrade(device)

    ##-----------------------------------------------------------------------------------##
    def shutdown(self, device):
    ##-----------------------------------------------------------------------------------##
        self.connect()
        self.client.exec_command('poweroff')

    ##-----------------------------------------------------------------------------------##
    def reboot(self, device):
    ##-----------------------------------------------------------------------------------##
        self.connect()
        self.client.exec_command('reboot')

    ##-----------------------------------------------------------------------------------##
    def upgrade(self, device):
    ##-----------------------------------------------------------------------------------##
        self.connect()
        self.client.exec_command('apt update && apt upgrade -y')