Commits

Enis Afgan committed 1628cf9

Add support for password-based instance login and FreeNX; clean up the code

  • Participants
  • Parent commits a8d93ff

Comments (0)

Files changed (1)

 import os, sys, yaml, urllib2, logging, hashlib, time, subprocess, random
 from urlparse import urlparse
 
+from boto.s3.key import Key
 from boto.s3.connection import S3Connection
-from boto.s3.key import Key
 from boto.exception import S3ResponseError
+
 logging.getLogger('boto').setLevel(logging.INFO) # Only log boto messages >=INFO
-
+log = None
 
 USER_DATA_URL = 'http://169.254.169.254/latest/user-data'
 # USER_DATA_URL = 'http://userwww.service.emory.edu/~eafgan/content/userData.yaml.sample' # used for testing
 # USER_DATA_URL = 'http://userwww.service.emory.edu/~eafgan/content/url_ud.txt' # used for testing
 LOCAL_PATH = '/tmp/cm' # Local path destination used for storing/reading any files created by this script
 USER_DATA_FILE_NAME = 'userData.yaml' # Local file with user data formatted by this script
-USER_DATA_FILE = os.path.join(LOCAL_PATH, USER_DATA_FILE_NAME) 
-USER_DATA_ORIG = os.path.join(LOCAL_PATH, 'original_%s' % USER_DATA_FILE_NAME) # Local file containing user data in its original format
+USER_DATA_FILE = os.path.join(LOCAL_PATH, USER_DATA_FILE_NAME) # The final/processed UD file
+# Local file containing UD in its original format
+USER_DATA_ORIG = os.path.join(LOCAL_PATH, 'original_%s' % USER_DATA_FILE_NAME)
 SERVICE_ROOT = 'http://s3.amazonaws.com/' # Obviously, customized for Amazon's S3
 DEFAULT_BUCKET_NAME = 'cloudman' # Ensure this bucket is accessible to anyone!
 DEFAULT_BOOT_SCRIPT_NAME = 'cm_boot.py' # Ensure this file is accessible to anyone in the public bucket!
 
 # ====================== Utility methods ======================
 
+def setup_logging():
+    # Logging setup
+    formatter = logging.Formatter("[%(levelname)s] %(module)s:%(lineno)d %(asctime)s: %(message)s")
+    console = logging.StreamHandler() # log to console - used during testing
+    # console.setLevel(logging.INFO) # accepts >INFO levels
+    console.setFormatter(formatter)
+    log_file = logging.FileHandler(os.path.join(LOCAL_PATH, "%s.log" % os.path.splitext(sys.argv[0])[0]), 'w')
+    log_file.setLevel(logging.DEBUG) # accepts all levels
+    log_file.setFormatter(formatter)
+    log = logging.root
+    log.addHandler(console)
+    log.addHandler(log_file)
+    log.setLevel(logging.DEBUG)
+    return log
+
 def _get_user_data():
+    ud = ''
     for i in range(0, 5):
           try:
               log.info("Getting user data from '%s', attempt %s" % (USER_DATA_URL, i))
     # return "gc_dev1|<account_key>|<secret_key>|somePWD"
     # with open('sample.yaml') as ud_yaml:
     #     ud = ud_yaml.read()
-    # return ud
+    if ud == '':
+        log.debug("Received empty/no user data")
+    return ud
 
 def _get_bucket_name(cluster_name, access_key):
     """Compose bucket name based on the user-provided cluster name and user access key""" 
 def _bucket_exists(s3_conn, bucket_name):
     bucket = None
     for i in range(1, 6):
-		try:
+        try:
             # log.debug("Looking for bucket '%s'" % bucket_name)
-		    bucket = s3_conn.lookup(bucket_name)
-		    break
-		except S3ResponseError: 
-		    log.error ("Bucket '%s' not found, attempt %s/5" % (bucket_name, i+1))
-		    time.sleep(2)
-		    
+            bucket = s3_conn.lookup(bucket_name)
+            break
+        except S3ResponseError: 
+            log.error ("Bucket '%s' not found, attempt %s/5" % (bucket_name, i+1))
+            time.sleep(2)
+            
     if bucket is not None:
         log.debug("Cluster bucket '%s' found." % bucket_name)
         return True
 def _remote_file_exists(s3_conn, bucket_name, remote_filename):
     b = None
     for i in range(0, 5):
-		try:
-			b = s3_conn.get_bucket(bucket_name)
-			break
-		except S3ResponseError: 
-			log.error ("Problem connecting to bucket '%s', attempt %s/5" % (bucket_name, i))
-			time.sleep(2)
-	    	
+        try:
+            b = s3_conn.get_bucket(bucket_name)
+            break
+        except S3ResponseError: 
+            log.error ("Problem connecting to bucket '%s', attempt %s/5" % (bucket_name, i))
+            time.sleep(2)
+            
     if b is not None:
-		k = Key(b, remote_filename)
-		if k.exists():
-		    return True
+        k = Key(b, remote_filename)
+        if k.exists():
+            return True
     return False
 
 def _save_file_to_bucket(s3_conn, bucket_name, remote_filename, local_file, force=False):
     # log.debug( "Establishing handle with bucket '%s'..." % bucket_name)
     b = None
     for i in range(0, 5):
-		try:
-			b = s3_conn.get_bucket(bucket_name)
-			break
-		except S3ResponseError, e:
-			log.error ("Problem connecting to bucket '%s', attempt %s/5" % (bucket_name, i))
-			time.sleep(2)
+        try:
+            b = s3_conn.get_bucket(bucket_name)
+            break
+        except S3ResponseError, e:
+            log.error ("Problem connecting to bucket '%s', attempt %s/5" % (bucket_name, i))
+            time.sleep(2)
     
     if b is not None:
         # log.debug("Establishing handle with key object '%s'..." % remote_filename)
     if ud.has_key('bucket_cluster') and ud['access_key'] is not None and ud['secret_key'] is not None:
         s3_conn = _get_s3_conn(ud['access_key'], ud['secret_key'])
         # Check if cluster bucket exists or use the default one
-        if not _bucket_exists(s3_conn, ud['bucket_cluster']) or not _remote_file_exists(s3_conn, ud['bucket_cluster'], ud['boot_script_name']):
+        if not _bucket_exists(s3_conn, ud['bucket_cluster']) or \
+           not _remote_file_exists(s3_conn, ud['bucket_cluster'], ud['boot_script_name']):
             log.debug("Using default bucket '%s'" % default_bucket_name)
             use_default_bucket = True
         else:
     if use_default_bucket is False:
         log.debug("Trying to get boot script '%s' from cluster bucket '%s'"
             % (ud['boot_script_name'], ud.get('bucket_cluster', None)))
-        got_boot_script = _get_file_from_bucket(s3_conn, ud['bucket_cluster'], ud['boot_script_name'], DEFAULT_BOOT_SCRIPT_NAME)
+        got_boot_script = _get_file_from_bucket(s3_conn, ud['bucket_cluster'], ud['boot_script_name'],
+            DEFAULT_BOOT_SCRIPT_NAME)
         if got_boot_script:
             os.chmod(os.path.join(LOCAL_PATH, DEFAULT_BOOT_SCRIPT_NAME), 0744)
     # If did not get boot script, fall back on the publicly available one
     if not got_boot_script or use_default_bucket:
-        boot_script_url = os.path.join(_get_default_bucket_url(ud), ud.get('boot_script_name', DEFAULT_BOOT_SCRIPT_NAME))
-        log.debug("Could not get boot script '%s' from cluster bucket '%s'; retrieving the public one from bucket url '%s'" 
+        boot_script_url = os.path.join(_get_default_bucket_url(ud), ud.get('boot_script_name', 
+            DEFAULT_BOOT_SCRIPT_NAME))
+        log.debug("Could not get boot script '%s' from cluster bucket '%s'; "
+            "retrieving the public one from bucket url '%s'" \
             % (ud['boot_script_name'], ud.get('bucket_cluster', None), boot_script_url))
         got_boot_script = _get_file_from_url(boot_script_url)
     if got_boot_script:
         # Save downloaded boot script to cluster bucket for future invocations
         if ud.has_key('bucket_cluster') and ud['bucket_cluster']:
             s3_conn = _get_s3_conn(ud['access_key'], ud['secret_key'])
-            if _bucket_exists(s3_conn, ud['bucket_cluster']) and not _remote_file_exists(s3_conn, ud['bucket_cluster'], ud['boot_script_name']):
-                _save_file_to_bucket(s3_conn, ud['bucket_cluster'], ud['boot_script_name'], DEFAULT_BOOT_SCRIPT_NAME)        
+            if _bucket_exists(s3_conn, ud['bucket_cluster']) and \
+               not _remote_file_exists(s3_conn, ud['bucket_cluster'], ud['boot_script_name']):
+                _save_file_to_bucket(s3_conn, ud['bucket_cluster'], ud['boot_script_name'], \
+                    DEFAULT_BOOT_SCRIPT_NAME)        
         return True
     log.debug("**Could not get the boot script**")
     return False
     # TODO: Check if th bucket 'default_bucket_name' is accessible to everyone 
     # because it is being accessed as a URL
     bucket_url = os.path.join(SERVICE_ROOT, default_bucket_name)
-    log.debug("Default bucket url: %s" % default_bucket_name)
+    log.debug("Default bucket url: %s" % bucket_url)
     return bucket_url
 
 def _handle_freenx(passwd):
     user = "ubuntu"
+    log.info("Setting up password-based login for user '{0}'".format(user))
     p1 = subprocess.Popen(["echo", "%s:%s" % (user, passwd)], stdout=subprocess.PIPE)
     p2 = subprocess.Popen(["chpasswd"], stdin=p1.stdout, stdout=subprocess.PIPE)
     p1.stdout.close()
     cl = ["sed", "-i", "s/^PasswordAuthentication .*/PasswordAuthentication yes/",
           "/etc/ssh/sshd_config"]
     subprocess.check_call(cl)
-    cl = ["/etc/init.d/ssh", "reload"]
+    cl = ["/usr/sbin/service", "ssh", "reload"]
     subprocess.check_call(cl)
-    cl = ["dpkg-reconfigure", "-pcritical", "freenx-server"]
-    # On slower/small instance types, there can be a conflict when running
-    # debconf so try this a few times
-    for i in range(5):
-        retcode = subprocess.call(cl)
-        if retcode == 0:
-            break
-        else:
-            time.sleep(5)
+    # Check if FreeNX is installed on the image before trying to configure it
+    cl = "/usr/bin/dpkg --get-selections | /bin/grep freenx"
+    retcode = subprocess.call(cl, shell=True)
+    if retcode == 0:
+        log.info("Setting up FreeNX")
+        cl = ["dpkg-reconfigure", "-pcritical", "freenx-server"]
+        # On slower/small instance types, there can be a conflict when running
+        # debconf so try this a few times
+        for i in range(5):
+            retcode = subprocess.call(cl)
+            if retcode == 0:
+                break
+            else:
+                time.sleep(5)
+    else:
+        log.info("freenx-server is not installed; not configuring it")
 
 # ====================== Actions methods ======================
 
         return
     if not ud.has_key('cluster_name'):
         log.warning("The provided user data should contain cluster_name field.")
-        ud['cluster_name'] = 'aGalaxyCloudManCluster_%s' % random.randrange(1, 9999999)
+        ud['cluster_name'] = 'aCloudManCluster_%s' % random.randrange(1, 9999999)
     elif ud['cluster_name'] == '':
         log.warning("The cluster_name field of user data should not be empty.")
-        ud['cluster_name'] = 'aGalaxyCloudManCluster_%s' % random.randrange(1, 9999999)
+        ud['cluster_name'] = 'aCloudManCluster_%s' % random.randrange(1, 9999999)
     
     if not ud.has_key('access_key'):
         log.info("The provided user data does not contain access_key field; setting it to None..")
     if _get_boot_script(ud):
         _run_boot_script(DEFAULT_BOOT_SCRIPT_NAME)
 
+# ====================== Driver code ======================
+
 def _parse_user_data(ud):
     if ud == '':
         _handle_empty()
     else: # default to yaml
         _handle_yaml(ud)
 
-# ====================== Driver code ======================
+def main():
+    if not os.path.exists(LOCAL_PATH):
+        os.mkdir(LOCAL_PATH)
+    global log
+    log = setup_logging()
+    ud = _get_user_data()
+    _parse_user_data(ud)
+    log.info("---> %s done <---" % sys.argv[0])
 
-if not os.path.exists(LOCAL_PATH):
-    os.mkdir(LOCAL_PATH)
-
-# Logging setup
-formatter = logging.Formatter("[%(levelname)s] %(module)s:%(lineno)d %(asctime)s: %(message)s")
-console = logging.StreamHandler() # log to console - used during testing
-# console.setLevel(logging.INFO) # accepts >INFO levels
-console.setFormatter(formatter)
-log_file = logging.FileHandler(os.path.join(LOCAL_PATH, "%s.log" % os.path.splitext(sys.argv[0])[0]), 'w') # log to a file
-log_file.setLevel(logging.DEBUG) # accepts all levels
-log_file.setFormatter(formatter)
-log = logging.root
-log.addHandler( console )
-log.addHandler( log_file )
-log.setLevel( logging.DEBUG )
-
-ud = _get_user_data()
-_parse_user_data(ud)
-log.info("---> %s done <---" % sys.argv[0])
+if __name__ == "__main__":
+    main()