refactor agent.py

context.py

+import platform
+
+""" This is the defautl context file. It replaces the Context class
+    to the platform specific one.
+"""
+
+system = platform.system()
+
+if system == "Windows":
+    from windows._win32context import Context
+    from win32.win32virtio import SerialPort
+elif system == "Linux":
+    from linux._linuxcontext import Context
+    from linux.posixvirtio import SerialPort
+
+else:
+    raise NotImplementedError("Platform %s is not supported.", system)
+
+
+class BaseContext():
+    pass
+
+Context
+SerialPort

linux/_linuxcontext.py

+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from os import mkdir, environ, chdir
+import platform
+from shutil import copy, rmtree, move
+import subprocess
+import sys
+
+system = platform.system()
+working_directory = sys.path[0]
+
+try:
+    chdir(working_directory)
+    subprocess.call(('pip', 'install', '-r', 'requirements.txt'))
+    if system == 'Linux':
+        copy("/root/agent/misc/vm_renewal", "/usr/local/bin/")
+except:
+    pass  # hope it works
+
+
+import logging
+import fileinput
+import tarfile
+from os.path import expanduser, join, exists
+from glob import glob
+from StringIO import StringIO
+from base64 import decodestring
+from hashlib import md5
+
+
+from ssh import PubKey
+from network import change_ip_ubuntu, change_ip_rhel
+
+from twisted.internet import reactor
+
+logging.basicConfig()
+logger = logging.getLogger()
+level = environ.get('LOGLEVEL', 'INFO')
+logger.setLevel(level)
+
+SSH_DIR = expanduser('~cloud/.ssh')
+AUTHORIZED_KEYS = join(SSH_DIR, 'authorized_keys')
+
+STORE_DIR = '/store'
+
+mount_template_linux = (
+    '//%(host)s/%(username)s %(dir)s cifs username=%(username)s'
+    ',password=%(password)s,iocharset=utf8,uid=cloud  0  0\n')
+
+
+distros = {'Scientific Linux': 'rhel',
+           'CentOS': 'rhel',
+           'CentOS Linux': 'rhel',
+           'Debian': 'debian',
+           'Ubuntu': 'debian'}
+if system == 'Linux':
+    distro = distros[platform.linux_distribution()[0]]
+
+
+class Context(object):
+
+    # http://stackoverflow.com/questions/12081310/
+    # python-module-to-change-system-date-and-time
+    def _linux_set_time(time):
+        import ctypes
+        import ctypes.util
+
+        CLOCK_REALTIME = 0
+
+        class timespec(ctypes.Structure):
+            _fields_ = [("tv_sec", ctypes.c_long),
+                        ("tv_nsec", ctypes.c_long)]
+
+        librt = ctypes.CDLL(ctypes.util.find_library("rt"))
+
+        ts = timespec()
+        ts.tv_sec = int(time)
+        ts.tv_nsec = 0
+
+        librt.clock_settime(CLOCK_REALTIME, ctypes.byref(ts))
+
+    @staticmethod
+    def change_password(password):
+        proc = subprocess.Popen(['/usr/sbin/chpasswd'],
+                                stdin=subprocess.PIPE)
+        proc.communicate('cloud:%s\n' % password)
+
+    @staticmethod
+    def restart_networking():
+        if distro == 'debian':
+            subprocess.call(['/etc/init.d/networking', 'restart'])
+        elif distro == 'rhel':
+            subprocess.call(['/bin/systemctl', 'restart', 'network'])
+            pass
+
+    @staticmethod
+    def change_ip(interfaces, dns):
+        if distro == 'debian':
+            change_ip_ubuntu(interfaces, dns)
+        elif distro == 'rhel':
+            change_ip_rhel(interfaces, dns)
+
+    @staticmethod
+    def set_time(time):
+        Context._linux_set_time(float(time))
+        try:
+            subprocess.call(['/etc/init.d/ntp', 'restart'])
+        except:
+            pass
+
+    @staticmethod
+    def set_hostname(hostname):
+        if distro == 'debian':
+            with open('/etc/hostname', 'w') as f:
+                f.write(hostname)
+        elif distro == 'rhel':
+            for line in fileinput.input('/etc/sysconfig/network',
+                                        inplace=1):
+                if line.startswith('HOSTNAME='):
+                    print 'HOSTNAME=%s' % hostname
+                else:
+                    print line.rstrip()
+
+        with open('/etc/hosts', 'w') as f:
+            f.write("127.0.0.1 localhost\n"
+                    "127.0.1.1 %s\n" % hostname)
+
+        subprocess.call(['/bin/hostname', hostname])
+
+    @staticmethod
+    def mount_store(host, username, password):
+        data = {'host': host, 'username': username, 'password': password}
+        data['dir'] = STORE_DIR
+        if not exists(STORE_DIR):
+            mkdir(STORE_DIR)
+        # TODO
+        for line in fileinput.input('/etc/fstab', inplace=True):
+            if not (line.startswith('//') and ' cifs ' in line):
+                print line.rstrip()
+
+        with open('/etc/fstab', 'a') as f:
+            f.write(mount_template_linux % data)
+
+        subprocess.call('mount -a', shell=True)
+
+    @staticmethod
+    def get_keys():
+        retval = []
+        try:
+            with open(AUTHORIZED_KEYS, 'r') as f:
+                for line in f.readlines():
+                    try:
+                        retval.append(PubKey.from_str(line))
+                    except:
+                        logger.exception(u'Invalid ssh key: ')
+        except IOError:
+            pass
+        return retval
+
+    @staticmethod
+    def _save_keys(keys):
+        print keys
+        try:
+            mkdir(SSH_DIR)
+        except OSError:
+            pass
+        with open(AUTHORIZED_KEYS, 'w') as f:
+            for key in keys:
+                f.write(unicode(key) + '\n')
+
+    @staticmethod
+    def add_keys(keys):
+        new_keys = Context.get_keys()
+        for key in keys:
+            try:
+                p = PubKey.from_str(key)
+                if p not in new_keys:
+                    new_keys.append(p)
+            except:
+                logger.exception(u'Invalid ssh key: ')
+        Context._save_keys(new_keys)
+
+    @staticmethod
+    def del_keys(keys):
+        new_keys = Context.get_keys()
+        for key in keys:
+            try:
+                p = PubKey.from_str(key)
+                try:
+                    new_keys.remove(p)
+                except ValueError:
+                    pass
+            except:
+                logger.exception(u'Invalid ssh key: ')
+        Context._save_keys(new_keys)
+
+    @staticmethod
+    def cleanup():
+        filelist = ([
+            '/root/.bash_history'
+            '/home/cloud/.bash_history'
+            '/root/.ssh'
+            '/home/cloud/.ssh']
+            + glob('/etc/ssh/ssh_host_*'))
+        for f in filelist:
+            rmtree(f, ignore_errors=True)
+
+        subprocess.call(('/usr/bin/ssh-keygen', '-A'))
+
+    @staticmethod
+    def start_access_server():
+        try:
+            subprocess.call(('/sbin/start', 'ssh'))
+        except OSError:
+            subprocess.call(('/bin/systemctl', 'start', 'sshd.service'))
+
+    @staticmethod
+    def append(data, filename, chunk_number, uuid):
+        if chunk_number == 0:
+            flag = "w"
+        else:
+            flag = "a"
+        with open(filename, flag) as myfile:
+            myfile.write(data)
+
+    @staticmethod
+    def update(filename, executable, checksum, uuid):
+        new_dir = working_directory + '.new'
+        old_dir = working_directory + '.old'
+        with open(filename, "r") as f:
+            data = f.read()
+            local_checksum = md5(data).hexdigest()
+            if local_checksum != checksum:
+                raise Exception("Checksum missmatch the file is damaged.")
+            decoded = StringIO(decodestring(data))
+        try:
+            tar = tarfile.TarFile.open("dummy", fileobj=decoded, mode='r|gz')
+            tar.extractall(new_dir)
+        except tarfile.ReadError as e:
+            logger.error(e)
+        move(working_directory, old_dir)
+        move(new_dir, working_directory)
+        logger.info("Transfer completed!")
+        reactor.stop()
+
+    @staticmethod
+    def ipaddresses():
+        import netifaces
+        args = {}
+        interfaces = netifaces.interfaces()
+        for i in interfaces:
+            if i == 'lo':
+                continue
+            args[i] = []
+            addresses = netifaces.ifaddresses(i)
+            args[i] = ([x['addr']
+                        for x in addresses.get(netifaces.AF_INET, [])] +
+                       [x['addr']
+                        for x in addresses.get(netifaces.AF_INET6, [])
+                        if '%' not in x['addr']])
+        return args
+
+    @staticmethod
+    def get_agent_version():
+        try:
+            with open('version.txt') as f:
+                return f.readline()
+        except IOError:
+            return None
+
+    @staticmethod
+    def send_expiration(url):
+        import notify
+        notify.notify(url)

linux/network.py

+import netifaces
+from netaddr import IPNetwork
+import fileinput
+import logging
+
+logger = logging.getLogger()
+
+interfaces_file = '/etc/network/interfaces'
+ifcfg_template = '/etc/sysconfig/network-scripts/ifcfg-%s'
+
+
+def get_interfaces_linux(interfaces):
+    for ifname in netifaces.interfaces():
+        mac = netifaces.ifaddresses(ifname)[17][0]['addr']
+        conf = interfaces.get(mac.upper())
+        if conf:
+            yield ifname, conf
+
+
+def remove_interfaces_ubuntu(devices):
+    delete_device = False
+
+    for line in fileinput.input(interfaces_file, inplace=True):
+        line = line.rstrip()
+        words = line.split()
+
+        if line.startswith('#') or line == '' or line.isspace() or not words:
+            # keep line
+            print line
+            continue
+
+        if (words[0] in ('auto', 'allow-hotplug') and
+                words[1].split(':')[0] in devices):
+            # remove line
+            continue
+
+        if words[0] == 'iface':
+            if words[1].split(':')[0] in devices:
+                # remove line
+                delete_device = True
+                continue
+            else:
+                delete_device = False
+
+        if line[0] in (' ', '\t') and delete_device:
+            # remove line
+            continue
+
+        # keep line
+        print line
+
+
+def change_ip_ubuntu(interfaces, dns):
+    data = list(get_interfaces_linux(interfaces))
+    remove_interfaces_ubuntu(dict(data).keys())
+
+    with open(interfaces_file, 'a') as f:
+        for ifname, conf in data:
+            ipv4_alias_counter = ipv6_alias_counter = 0
+            f.write('auto %s\n' % ifname)
+            for i in conf['addresses']:
+                ip_with_prefix = IPNetwork(i)
+                prefixlen = ip_with_prefix.prefixlen
+                ip = ip_with_prefix.ip
+                alias = ifname
+                if ip.version == 6:
+                    if ipv6_alias_counter > 0:
+                        alias = '%s:%d' % (ifname, ipv6_alias_counter)
+                    ipv6_alias_counter += 1
+                else:
+                    if ipv4_alias_counter > 0:
+                        alias = '%s:%d' % (ifname, ipv4_alias_counter)
+                    ipv4_alias_counter += 1
+
+                f.write(
+                    'iface %(ifname)s %(proto)s static\n'
+                    '    address %(ip)s\n'
+                    '    netmask %(prefixlen)d\n'
+                    '    gateway %(gw)s\n'
+                    '    dns-nameservers %(dns)s\n' % {
+                        'ifname': alias,
+                        'proto': 'inet6' if ip.version == 6 else 'inet',
+                        'ip': ip,
+                        'prefixlen': prefixlen,
+                        'gw': conf['gw6' if ip.version == 6 else 'gw4'],
+                        'dns': dns})
+
+
+# example:
+# change_ip_ubuntu({
+#    u'02:00:00:02:A3:E8': {
+#        u'gw4': u'10.1.0.254', 'gw6': '2001::ffff',
+#        u'addresses': [u'10.1.0.84/24', '10.1.0.1/24', '2001::1/48']},
+#    u'02:00:00:02:A3:E9': {
+#        u'gw4': u'10.255.255.1', u'addresses': [u'10.255.255.9']}},
+#    '8.8.8.8')
+
+
+def change_ip_rhel(interfaces, dns):
+    for ifname, conf in get_interfaces_linux(interfaces):
+        with open(ifcfg_template % ifname,
+                  'w') as f:
+            f.write('DEVICE=%s\n'
+                    'BOOTPROTO=none\n'
+                    'USERCTL=no\n'
+                    'ONBOOT=yes\n' % ifname)
+            for i in conf['addresses']:
+                ip_with_prefix = IPNetwork(i)
+                ip = ip_with_prefix.ip
+                if ip.version == 6:
+                    f.write('IPV6INIT=yes\n'
+                            'IPV6ADDR=%(ip)s/%(prefixlen)d\n'
+                            'IPV6_DEFAULTGW=%(gw)s\n' % {
+                                'ip': ip,
+                                'prefixlen': ip_with_prefix.prefixlen,
+                                'gw': conf['gw6']})
+                else:
+                    f.write('NETMASK=%(netmask)s\n'
+                            'IPADDR=%(ip)s\n'
+                            'GATEWAY=%(gw)s\n' % {
+                                'ip': ip,
+                                'netmask': str(ip_with_prefix.netmask),
+                                'gw': conf['gw4']})

linux/posixvirtio.py

+# Copyright (c) Twisted Matrix Laboratories.
+# See LICENSE for details.
+
+
+"""
+Virtio-Serial Port Protocol
+"""
+
+# system imports
+import os
+
+# dependent on pyserial ( http://pyserial.sf.net/ )
+# only tested w/ 1.18 (5 Dec 2002)
+
+# twisted imports
+from twisted.internet import abstract, fdesc
+
+
+class SerialPort(abstract.FileDescriptor):
+    """
+    A select()able serial device, acting as a transport.
+    """
+
+    connected = 1
+
+    def __init__(self, protocol, deviceNameOrPortNumber, reactor):
+        abstract.FileDescriptor.__init__(self, reactor)
+        self.port = deviceNameOrPortNumber
+        self._serial = os.open(
+            self.port, os.O_RDWR | os.O_NOCTTY | os.O_NONBLOCK)
+        self.reactor = reactor
+        self.protocol = protocol
+        self.protocol.makeConnection(self)
+        self.startReading()
+
+    def fileno(self):
+        return self._serial
+
+    def writeSomeData(self, data):
+        """
+        Write some data to the serial device.
+        """
+        return fdesc.writeToFD(self.fileno(), data)
+
+    def doRead(self):
+        """
+        Some data's readable from serial device.
+        """
+        return fdesc.readFromFD(self.fileno(), self.protocol.dataReceived)
+
+    def connectionLost(self, reason):
+        """
+        Called when the serial port disconnects.
+
+        Will call C{connectionLost} on the protocol that is handling the
+        serial data.
+        """
+        abstract.FileDescriptor.connectionLost(self, reason)
+        os.close(self._serial)
+        self.protocol.connectionLost(reason)