Support openssh public key format

azurelinuxagent/distro/default/deprovision.py

@@ -50,7 +50,7 @@ class DeprovisionHandler(object):
             warnings.append("WARNING! Skip delete user.")
             return
 
-        username = ovfenv.get_username()
+        username = ovfenv.username
         warnings.append(("WARNING! {0} account and entire home directory "
                          "will be deleted.").format(username))
         actions.append(DeprovisionAction(OSUTIL.del_account, [username]))

azurelinuxagent/distro/default/osutil.py

@@ -78,24 +78,6 @@ class DefaultOSUtil(object):
     def get_openssl_cmd(self):
         return self.openssl_cmd
 
-    def set_user_account(self, username, password, expiration=None):
-        """
-        Update password and ssh key for user account.
-        New account will be created if not exists.
-        """
-        if username is None:
-            raise OSUtilError("User name is empty")
-
-        if self.is_sys_user(username):
-            raise OSUtilError(("User {0} is a system user. "
-                               "Will not set passwd.").format(username))
-
-        userentry = self.get_userentry(username)
-        if userentry is None:
-            self.useradd(username, expiration)
-
-        self.conf_sudoer(username, password is None)
-
     def get_userentry(self, username):
         try:
             return pwd.getpwnam(username)
@@ -106,7 +88,8 @@ class DefaultOSUtil(object):
         userentry = self.get_userentry(username)
         uidmin = None
         try:
-            uidmin_def = fileutil.get_line_startingwith("UID_MIN", "/etc/login.defs")
+            uidmin_def = fileutil.get_line_startingwith("UID_MIN", 
+                                                        "/etc/login.defs")
             if uidmin_def is not None:
                 uidmin = int(uidmin_def.split()[1])
         except IOError as e:
@@ -119,6 +102,10 @@ class DefaultOSUtil(object):
             return False
 
     def useradd(self, username, expiration=None):
+        """
+        Update password and ssh key for user account.
+        New account will be created if not exists.
+        """
         if expiration is not None:
             cmd = "useradd -m {0} -e {1}".format(username, expiration)
         else:
@@ -130,7 +117,10 @@ class DefaultOSUtil(object):
                                "output:{2}").format(username, retcode, out))
 
     def chpasswd(self, username, password, use_salt=True, salt_type=6,
-                       salt_len=10):
+                 salt_len=10):
+        if self.is_sys_user(username):
+            raise OSUtilError(("User {0} is a system user. "
+                               "Will not set passwd.").format(username))
         passwd_hash = textutil.gen_password_hash(password, use_salt, salt_type,
                                                  salt_len)
         try:
@@ -197,19 +187,18 @@ class DefaultOSUtil(object):
         thumbprint = thumbprint.rstrip().split('=')[1].replace(':', '').upper()
         return thumbprint
 
-    def deploy_ssh_keypair(self, username, thumbprint, path):
+    def deploy_ssh_keypair(self, username, keypair):
         """
         Deploy id_rsa and id_rsa.pub
         """
+        path, thumbprint = keypair
         path = self._norm_path(path)
         dir_path = os.path.dirname(path)
         fileutil.mkdir(dir_path, mode=0700, owner=username)
         lib_dir = self.get_lib_dir()
         prv_path = os.path.join(lib_dir, thumbprint + '.prv')
         if not os.path.isfile(prv_path):
-            logger.error("Failed to deploy key pair, thumbprint: {0}",
-                         thumbprint)
-            return
+            raise OSUtilError("Can't find {0}.prv".format(thumbprint))
         shutil.copyfile(prv_path, path)
         pub_path = path + '.pub'
         pub = self.get_pubkey_from_prv(prv_path)
@@ -223,28 +212,39 @@ class DefaultOSUtil(object):
         shellutil.run("ssh-keygen -i -m PKCS8 -f {0} >> {1}".format(input_file,
                                                                     output_file))
 
-    def deploy_ssh_pubkey(self, username, thumbprint, path):
+    def deploy_ssh_pubkey(self, username, pubkey):
         """
         Deploy authorized_key
         """
+        path, thumbprint, value = pubkey
+        if path is None:
+            raise OSUtilError("Publich key path is None")
+
         path = self._norm_path(path)
         dir_path = os.path.dirname(path)
         fileutil.mkdir(dir_path, mode=0700, owner=username)
-        lib_dir = self.get_lib_dir()
-        crt_path = os.path.join(lib_dir, thumbprint + '.crt')
-        if not os.path.isfile(crt_path):
-            logger.error("Failed to deploy public key, thumbprint: {0}",
-                         thumbprint)
-            return
-        pub_path = os.path.join(lib_dir, thumbprint + '.pub')
-        pub = self.get_pubkey_from_crt(crt_path)
-        fileutil.write_file(pub_path, pub)
-        self.set_selinux_context(pub_path, 'unconfined_u:object_r:ssh_home_t:s0')
-        self.openssl_to_openssh(pub_path, path)
+        if value is not None:
+            if not value.startswith("ssh-"):
+                raise OSUtilError("Bad public key: {0}".format(value))
+            fileutil.write_file(path, value)
+        elif thumbprint is not None:
+            lib_dir = self.get_lib_dir()
+            crt_path = os.path.join(lib_dir, thumbprint + '.crt')
+            if not os.path.isfile(crt_path):
+                raise OSUtilError("Can't find {0}.crt".format(thumbprint))
+            pub_path = os.path.join(lib_dir, thumbprint + '.pub')
+            pub = self.get_pubkey_from_crt(crt_path)
+            fileutil.write_file(pub_path, pub)
+            self.set_selinux_context(pub_path, 
+                                     'unconfined_u:object_r:ssh_home_t:s0')
+            self.openssl_to_openssh(pub_path, path)
+            fileutil.chmod(pub_path, 0600)
+        else:
+            raise OSUtilError("SSH public key Fingerprint and Value are None")
+
         self.set_selinux_context(path, 'unconfined_u:object_r:ssh_home_t:s0')
         fileutil.chowner(path, username)
         fileutil.chmod(path, 0644)
-        fileutil.chmod(pub_path, 0600)
 
     def is_selinux_system(self):
         """

azurelinuxagent/distro/default/provision.py

@@ -15,12 +15,16 @@
 # Requires Python 2.4+ and Openssl 1.0+
 #
 
+"""
+Provision handler
+"""
+
 import os
 import azurelinuxagent.logger as logger
 import azurelinuxagent.conf as conf
 from azurelinuxagent.event import add_event, WALAEventOperation
 from azurelinuxagent.exception import *
-from azurelinuxagent.utils.osutil import OSUTIL
+from azurelinuxagent.utils.osutil import OSUTIL, OSUtilError
 import azurelinuxagent.protocol as prot
 import azurelinuxagent.protocol.ovfenv as ovf
 import azurelinuxagent.utils.shellutil as shellutil
@@ -66,7 +70,6 @@ class ProvisionHandler(object):
             add_event(name="WALA", is_success=False, message=str(e),
                               op=WALAEventOperation.Provision)
 
-
     def reg_ssh_host_key(self):
         keypair_type = conf.get("Provisioning.SshHostKeyPairType", "rsa")
         if conf.get_switch("Provisioning.RegenerateSshHostKeyPair"):
@@ -92,26 +95,40 @@ class ProvisionHandler(object):
             ovfenv = ovf.copy_ovf_env()
         except prot.ProtocolError as e:
             raise ProvisionError("Failed to copy ovf-env.xml: {0}".format(e))
+    
+        logger.info("Handle ovf-env.xml.")
+        try:
+            logger.info("Set host name.")
+            OSUTIL.set_hostname(ovfenv.hostname)
 
-        password = ovfenv.get_user_password()
-        ovfenv.clear_user_password()
+            logger.info("Publish host name.")
+            OSUTIL.publish_hostname(ovfenv.hostname)
 
-        logger.info("Set host name.")
-        OSUTIL.set_hostname(ovfenv.get_computer_name())
-        logger.info("Publish host name.")
-        OSUTIL.publish_hostname(ovfenv.get_computer_name())
-        logger.info("Create user account.")
-        OSUTIL.set_user_account(ovfenv.get_username(), password)
+            self.config_user_account(ovfenv)
 
-        if password is not None:
+            self.save_customdata(ovfenv)
+
+            if conf.get_switch("Provisioning.DeleteRootPassword"):
+                OSUTIL.del_root_password()
+        except OSUtilError as e:
+            raise ProvisionError("Failed to handle ovf-env.xml: {0}".format(e))
+        
+    def config_user_account(self, ovfenv):
+        logger.info("Create user account if not exists")
+        OSUTIL.useradd(ovfenv.username)
+
+        if ovfenv.user_password is not None:
+            logger.info("Set user password.")
             use_salt = conf.get_switch("Provision.UseSalt", True)
             salt_type = conf.get_switch("Provision.SaltType", 6)
-            logger.info("Set user password.")
-            OSUTIL.chpasswd(ovfenv.get_username(), password, use_salt,
-                            salt_type)
+            OSUTIL.chpasswd(ovfenv.username, ovfenv.user_password, 
+                            use_salt,salt_type)
+         
+        logger.info("Configure sudoer")
+        OSUTIL.conf_sudoer(ovfenv.username, ovfenv.user_password is None)
 
-        logger.info("Configure sshd.")
-        OSUTIL.conf_sshd(ovfenv.get_disable_ssh_password_auth())
+        logger.info("Configure sshd")
+        OSUTIL.conf_sshd(ovfenv.disable_ssh_password_auth)
 
         #Disable selinux temporary
         sel = OSUTIL.is_selinux_enforcing()
@@ -120,20 +137,15 @@ class ProvisionHandler(object):
 
         self.deploy_ssh_pubkeys(ovfenv)
         self.deploy_ssh_keypairs(ovfenv)
-        self.save_customdata(ovfenv)
 
         if sel:
             OSUTIL.set_selinux_enforce(1)
 
         OSUTIL.restart_ssh_service()
 
-        if conf.get_switch("Provisioning.DeleteRootPassword"):
-            OSUTIL.del_root_password()
-
-
     def save_customdata(self, ovfenv):
         logger.info("Save custom data")
-        customdata = ovfenv.get_customdata()
+        customdata = ovfenv.customdata
         if customdata is None:
             return
         lib_dir = OSUTIL.get_lib_dir()
@@ -141,12 +153,12 @@ class ProvisionHandler(object):
                             OSUTIL.decode_customdata(customdata))
 
     def deploy_ssh_pubkeys(self, ovfenv):
-        for thumbprint, path in ovfenv.get_ssh_pubkeys():
+        for pubkey in ovfenv.ssh_pubkeys:
             logger.info("Deploy ssh public key.")
-            OSUTIL.deploy_ssh_pubkey(ovfenv.get_username(), thumbprint, path)
+            OSUTIL.deploy_ssh_pubkey(ovfenv.username, pubkey)
 
     def deploy_ssh_keypairs(self, ovfenv):
-        for thumbprint, path in ovfenv.get_ssh_keypairs():
+        for keypair in ovfenv.ssh_keypairs:
             logger.info("Deploy ssh key pairs.")
-            OSUTIL.deploy_ssh_keypair(ovfenv.get_username(), thumbprint, path)
+            OSUTIL.deploy_ssh_keypair(ovfenv.username, keypair)
 

azurelinuxagent/protocol/ovfenv.py

@@ -16,12 +16,30 @@
 #
 # Requires Python 2.4+ and Openssl 1.0+
 #
-
-from azurelinuxagent.protocol.common import *
-
+"""
+Copy and parse ovf-env.xml from provisiong ISO and local cache
+"""
+import os
+import re
+import xml.etree.ElementTree as ET
+import azurelinuxagent.logger as logger
+import azurelinuxagent.utils.fileutil as fileutil
+from azurelinuxagent.utils.textutil import find_text
 from azurelinuxagent.utils.osutil import OSUTIL, OSUtilError
+from azurelinuxagent.protocol import ProtocolError
+
+OVF_FILE_NAME = "ovf-env.xml"
+OVF_VERSION = "1.0"
+OVF_NAME_SPACE = {
+        "oe" : "http://schemas.dmtf.org/ovf/environment/1",
+        "wa" : "http://schemas.microsoft.com/windowsazure",
+        "i" : "http://www.w3.org/2001/XMLSchema-instance"
+}
 
 def get_ovf_env():
+    """
+    Load saved ovf-env.xml
+    """
     ovf_file_path = os.path.join(OSUTIL.get_lib_dir(), OVF_FILE_NAME)
     if os.path.isfile(ovf_file_path):
         xml_text = fileutil.read_file(ovf_file_path)
@@ -37,7 +55,6 @@ def copy_ovf_env():
     try:
         OSUTIL.mount_dvd()
         ovf_file_path_on_dvd = OSUTIL.get_ovf_env_file_path_on_dvd()
-
         ovfxml = fileutil.read_file(ovf_file_path_on_dvd, remove_bom=True)
         ovfenv = OvfEnv(ovfxml)
         ovfxml = re.sub("<UserPassword>.*?<", "<UserPassword>*<", ovfxml)
@@ -50,7 +67,10 @@ def copy_ovf_env():
         raise ProtocolError(str(e))
     return ovfenv
 
-OVF_FILE_NAME="ovf-env.xml"
+def _validate_ovf(val, msg):
+    if val is None:
+        raise ProtocolError("Failed to parse OVF XML: {0}".format(msg))
+
 class OvfEnv(object):
     """
     Read, and process provisioning info from provisioning file OvfEnv.xml
@@ -59,121 +79,61 @@ class OvfEnv(object):
         if xml_text is None:
             raise ValueError("ovf-env is None")
         logger.verb("Load ovf-env.xml")
-        self.parse(xml_text)
-
-    def reinitialize(self):
-        """
-        Reset members.
-        """
-        self.wa_ns = "http://schemas.microsoft.com/windowsazure"
-        self.ovf_ns = "http://schemas.dmtf.org/ovf/environment/1"
-        self.major_version = 1
-        self.minor_version = 0
-        self.compute_name = None
-        self.user_name = None
+        self.hostname = None
+        self.username = None
         self.user_password = None
         self.customdata = None
         self.disable_ssh_password_auth = True
         self.ssh_pubkeys = []
         self.ssh_keypairs = []
-
-    def get_major_version(self):
-        return self.major_version
-
-    def get_minor_version(self):
-        return self.minor_version
-
-    def get_computer_name(self):
-        return self.compute_name
-
-    def get_username(self):
-        return self.user_name
-
-    def get_user_password(self):
-        return self.user_password
-
-    def clear_user_password(self):
-        self.user_password = None
-
-    def get_customdata(self):
-        return self.customdata
-
-    def get_disable_ssh_password_auth(self):
-        return self.disable_ssh_password_auth
-
-    def get_ssh_pubkeys(self):
-        return self.ssh_pubkeys
-
-    def get_ssh_keypairs(self):
-        return self.ssh_keypairs
+        self.parse(xml_text)
 
     def parse(self, xml_text):
         """
         Parse xml tree, retreiving user and ssh key information.
         Return self.
         """
-        self.reinitialize()
-        dom = xml.dom.minidom.parseString(xml_text)
-        if len(dom.getElementsByTagNameNS(self.ovf_ns, "Environment")) != 1:
-            logger.error("Unable to parse OVF XML.")
-        section = None
-        newer = False
-        for p in dom.getElementsByTagNameNS(self.wa_ns, "ProvisioningSection"):
-            for n in p.childNodes:
-                if n.localName == "Version":
-                    verparts = get_node_text(n).split('.')
-                    major = int(verparts[0])
-                    minor = int(verparts[1])
-                    if major > self.major_version:
-                        newer = True
-                    if major != self.major_version:
-                        break
-                    if minor > self.minor_version:
-                        newer = True
-                    section = p
-        if newer == True:
+        ns = OVF_NAME_SPACE
+        xml_doc = ET.fromstring(xml_text)
+        section = xml_doc.find(".//wa:ProvisioningSection", ns)
+        _validate_ovf(section, "ProvisioningSection not found")
+
+        version = section.find("wa:Version", ns).text
+        _validate_ovf(version, "Version not found")
+
+        if version > OVF_VERSION:
             logger.warn("Newer provisioning configuration detected. "
-                    "Please consider updating waagent.")
-            if section == None:
-                logger.error("Could not find ProvisioningSection with "
-                        "major version={0}", self.major_version)
-                return None
-        self.compute_name = get_node_text(section.getElementsByTagNameNS(self.wa_ns, "HostName")[0])
-        self.user_name = get_node_text(section.getElementsByTagNameNS(self.wa_ns, "UserName")[0])
-        try:
-            self.user_password = get_node_text(section.getElementsByTagNameNS(self.wa_ns, "UserPassword")[0])
-        except:
-            pass
-        cd_section=None
-        cd_section=section.getElementsByTagNameNS(self.wa_ns, "CustomData")
-        if len(cd_section) > 0 :
-            self.customdata=get_node_text(cd_section[0])
-        disable_ssh_password_auth = section.getElementsByTagNameNS(self.wa_ns, "DisableSshPasswordAuthentication")
-        if len(disable_ssh_password_auth) != 0:
-            self.disable_ssh_password_auth = (get_node_text(disable_ssh_password_auth[0]).lower() == "true")
-        for pkey in section.getElementsByTagNameNS(self.wa_ns, "PublicKey"):
-            logger.verb(repr(pkey))
-            fp = None
-            path = None
-            for c in pkey.childNodes:
-                if c.localName == "Fingerprint":
-                    fp = get_node_text(c).upper()
-                    logger.verb(fp)
-                if c.localName == "Path":
-                    path = get_node_text(c)
-                    logger.verb(path)
-            self.ssh_pubkeys += [[fp, path]]
-        for keyp in section.getElementsByTagNameNS(self.wa_ns, "KeyPair"):
-            fp = None
-            path = None
-            logger.verb(repr(keyp))
-            for c in keyp.childNodes:
-                if c.localName == "Fingerprint":
-                    fp = get_node_text(c).upper()
-                    logger.verb(fp)
-                if c.localName == "Path":
-                    path = get_node_text(c)
-                    logger.verb(path)
-            self.ssh_keypairs += [[fp, path]]
-        return self
+                        "Please consider updating waagent")
+
+        conf_set = section.find("wa:LinuxProvisioningConfigurationSet", ns)
+        _validate_ovf(conf_set, "LinuxProvisioningConfigurationSet not found")
+
+        self.hostname = find_text(conf_set, "wa:HostName", ns=ns)
+        _validate_ovf(self.hostname, "HostName not found")
+
+        self.username = find_text(conf_set, "wa:UserName", ns=ns)
+        _validate_ovf(self.username, "UserName not found")
+
+        self.user_password = find_text(conf_set, "wa:UserPassword", ns=ns)
+
+        self.customdata = find_text(conf_set, "wa:CustomData", ns=ns)
+
+        auth = find_text(conf_set, "wa:DisableSshPasswordAuthentication", ns=ns)
+        if auth is not None and auth.lower() == "true":
+            self.disable_ssh_password_auth = True
+        else:
+            self.disable_ssh_password_auth = False
+
+        public_keys = conf_set.findall("wa:SSH/wa:PublicKeys/wa:PublicKey", ns)
+        for public_key in public_keys:
+            path = find_text(public_key, "wa:Path", ns=ns)
+            fingerprint = find_text(public_key, "wa:Fingerprint", ns=ns)
+            value = find_text(public_key, "wa:Value", ns=ns)
+            self.ssh_pubkeys.append((path, fingerprint, value))
+
+        keypairs = conf_set.findall("wa:SSH/wa:KeyPairs/wa:KeyPair", ns)
+        for keypair in keypairs:
+            path = find_text(keypair, "wa:Path", ns=ns)
+            fingerprint = find_text(keypair, "wa:Fingerprint", ns=ns)
+            self.ssh_keypairs.append((path, fingerprint))
 

azurelinuxagent/utils/textutil.py

@@ -21,15 +21,22 @@ import random
 import string
 import struct
 
-def find_first_node(xml_doc, selector):
-    nodes = find_all_nodes(xml_doc, selector)
+def find_first_node(xml_doc, selector, ns=None):
+    nodes = find_all_nodes(xml_doc, selector, ns=ns)
     if len(nodes) > 0:
         return nodes[0]
 
-def find_all_nodes(xml_doc, selector):
-    nodes = xml_doc.findall(selector)
+def find_all_nodes(xml_doc, selector, ns=None):
+    nodes = xml_doc.findall(selector, ns)
     return nodes
 
+def find_text(root, selector, ns=None, default=None):
+    element = root.find(selector, ns)
+    if element is not None:
+        return element.text
+    else:
+        return default
+
 def get_node_text(a):
     """
     Filter non-text nodes from DOM tree

tests/test_osutil.py

@@ -28,7 +28,7 @@ import time
 import azurelinuxagent.utils.fileutil as fileutil
 import azurelinuxagent.utils.shellutil as shellutil
 import azurelinuxagent.conf as conf
-from azurelinuxagent.utils.osutil import OSUTIL
+from azurelinuxagent.utils.osutil import OSUTIL, OSUtilError
 import test
 
 class TestOSUtil(unittest.TestCase):
@@ -62,8 +62,9 @@ class TestCurrOS(unittest.TestCase):
     @mock(shellutil, 'run', MockFunc())
     @mock(shellutil, 'run_get_output', MockFunc(retval=[0, '']))
     def test_update_user_account(self):
-        OSUTIL.set_user_account('api', 'api')
-        OSUTIL.del_account('api')
+        OSUTIL.useradd('foo')
+        OSUTIL.chpasswd('foo', 'bar')
+        OSUTIL.del_account('foo')
 
     @mock(fileutil, 'read_file', MockFunc(retval='root::::'))
     @mock(fileutil, 'write_file', MockFunc())
@@ -119,12 +120,19 @@ class TestCurrOS(unittest.TestCase):
         OSUTIL.publish_hostname('api')
    
     @mock(OSUTIL, 'get_home', MockFunc(retval='/tmp/home'))
+    @mock(OSUTIL, 'get_pubkey_from_prv', MockFunc(retval=''))
+    @mock(fileutil, 'chowner', MockFunc())
     def test_deploy_key(self):
         if os.path.isdir('/tmp/home'):
             shutil.rmtree('/tmp/home')
-        user = shellutil.run_get_output('whoami')[1].strip()
-        OSUTIL.deploy_ssh_keypair(user, 'test', '$HOME/.ssh/id_rsa')
-        OSUTIL.deploy_ssh_pubkey(user, 'test', '$HOME/.ssh/authorized_keys')
+        fileutil.write_file('/tmp/foo.prv', '')
+        OSUTIL.deploy_ssh_keypair("foo", ('$HOME/.ssh/id_rsa', 'foo'))
+        OSUTIL.deploy_ssh_pubkey("foo", ('$HOME/.ssh/authorized_keys', None, 
+                                         'ssh-rsa asdf'))
+        OSUTIL.deploy_ssh_pubkey("foo", ('$HOME/.ssh/authorized_keys', 'foo', 
+                                         'ssh-rsa asdf'))
+        self.assertRaises(OSUtilError, OSUTIL.deploy_ssh_pubkey, "foo", 
+                         ('$HOME/.ssh/authorized_keys', 'foo','hehe-rsa asdf'))
         self.assertTrue(os.path.isfile('/tmp/home/.ssh/id_rsa'))
         self.assertTrue(os.path.isfile('/tmp/home/.ssh/id_rsa.pub'))
         self.assertTrue(os.path.isfile('/tmp/home/.ssh/authorized_keys'))

tests/test_ovfxml.py

@@ -41,6 +41,7 @@ ExtensionsConfigSample="""
             <PublicKey>
               <Fingerprint>EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62</Fingerprint>
               <Path>$HOME/UserName/.ssh/authorized_keys</Path>
+              <Value>ssh-rsa AAAANOTAREALKEY== foo@bar.local</Value>
             </PublicKey>
           </PublicKeys>
           <KeyPairs>
@@ -59,17 +60,20 @@ ExtensionsConfigSample="""
 class TestOvf(unittest.TestCase):
     def test_ovf(self):
         config = ovfenv.OvfEnv(ExtensionsConfigSample)
-        self.assertEquals(1, config.get_major_version())
-        self.assertEquals(0, config.get_minor_version())
-        self.assertEquals("HostName", config.get_computer_name())
-        self.assertEquals("UserName", config.get_username())
-        self.assertEquals("UserPassword", config.get_user_password())
-        self.assertEquals(False, config.get_disable_ssh_password_auth())
-        self.assertEquals("CustomData", config.get_customdata())
-        self.assertNotEquals(None, config.get_ssh_pubkeys())
-        self.assertEquals(1, len(config.get_ssh_pubkeys()))
-        self.assertNotEquals(None, config.get_ssh_keypairs())
-        self.assertEquals(1, len(config.get_ssh_keypairs()))
+        self.assertEquals("HostName", config.hostname)
+        self.assertEquals("UserName", config.username)
+        self.assertEquals("UserPassword", config.user_password)
+        self.assertEquals(False, config.disable_ssh_password_auth)
+        self.assertEquals("CustomData", config.customdata)
+        self.assertNotEquals(None, config.ssh_pubkeys)
+        self.assertEquals(1, len(config.ssh_pubkeys))
+        pubkey = config.ssh_pubkeys[0]
+        path, fingerprint, value = pubkey
+        self.assertEquals(path, "$HOME/UserName/.ssh/authorized_keys")
+        self.assertEquals(fingerprint, "EB0C0AB4B2D5FC35F2F0658D19F44C8283E2DD62"),
+        self.assertEquals(value, "ssh-rsa AAAANOTAREALKEY== foo@bar.local")
+        self.assertNotEquals(None, config.ssh_keypairs)
+        self.assertEquals(1, len(config.ssh_keypairs))
         
 if __name__ == '__main__':
     unittest.main()

tests/tools.py

@@ -21,6 +21,7 @@
 
 import os
 import sys
+from functools import wraps
 from azurelinuxagent.utils.osutil import OSUTIL
 
 parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
@@ -33,6 +34,7 @@ def simple_file_grep(file_path, search_str):
 
 def mock(target, name, mock):
     def decorator(func):
+        @wraps(func)
         def wrapper(*args, **kwargs):
             origin = getattr(target, name)
             setattr(target, name, mock)