#!/usr/bin/env python

from netmiko import ConnectHandler
import getpass
import textwrap
import getpass
from os.path import expanduser


class CiscoIOS:

    def __init__(self, device_param):
        self.device_param = device_param
        self.ssh = ConnectHandler(**device_param)

    def load_ssh_key(self, id_rsa=None, user=None):
        if not id_rsa:
            id_rsa = expanduser('~') + '/.ssh/id_rsa.pub'
        if not user:
            user = getpass.getuser()
        pubkey = open(id_rsa, 'r').read()
        pubkey = textwrap.wrap(pubkey, 80)

        commands = ['ip ssh pubkey-chain',
                    'username ' + user,
                    'key-string'] + pubkey + ['exit']

        return self.ssh.send_config_set(commands)

    def del_ssh_key(self, user):
        commands = ['ip ssh pubkey-chain',
                    'no username ' + user,
                    'exit']
        return self.ssh.send_config_set(commands)

    def write_memory(self):
        ret = 'enable password is not set'
        if not self.ssh.check_enable_mode() and \
                self.device_param['secret']:
            self.ssh.enable()
            ret = self.ssh.send_command('write memory')
            self.ssh.exit_enable_mode()
        return ret

if __name__ == '__main__':
    pwd = getpass.getpass()
    device = {'device_type': 'cisco_ios',
              'ip': 'acc-sw-1',
              'username': 'anon',
              'password': pwd,
              'secret': ''}
    cisco = CiscoIOS(device)
    print(cisco.load_ssh_key(user='anon'))
    print(cisco.write_memory())