merged master

utils/hosting_utils.py

@@ -0,0 +1,11 @@
+from hosting.models import UserHostingKey
+
+
+def get_all_public_keys(customer):
+    """
+    Returns all the public keys of the user
+    :param customer: The customer whose public keys are needed
+    :return: A list of public keys
+    """
+    return UserHostingKey.objects.filter(user_id=customer.id).values_list(
+        "public_key", flat=True)

utils/tasks.py

@@ -1,8 +1,15 @@
+import tempfile
+
+import cdist
+from cdist.integration import configure_hosts_simple
+from celery.result import AsyncResult
+from celery import current_task
 from celery.utils.log import get_task_logger
 from django.conf import settings
-from dynamicweb.celery import app
 from django.core.mail import EmailMessage
 
+from dynamicweb.celery import app
+
 logger = get_task_logger(__name__)
 
 
@@ -18,3 +25,74 @@ def send_plain_email_task(self, email_data):
     """
     email = EmailMessage(**email_data)
     email.send()
+
+
+@app.task(bind=True, max_retries=settings.CELERY_MAX_RETRIES)
+def save_ssh_key(self, hosts, keys):
+    """
+    Saves ssh key into the VMs of a user using cdist
+
+    :param hosts: A list of hosts to be configured
+    :param keys: A list of keys to be added. A key should be dict of the
+           form    {
+                       'value': 'sha-.....', # public key as string
+                       'state': True         # whether key is to be added or
+                    }                        # removed
+    """
+    logger.debug(
+        "Running save_ssh_key on {}".format(current_task.request.hostname))
+    logger.debug("""Running save_ssh_key task for
+                    Hosts: {hosts_str}
+                    Keys: {keys_str}""".format(hosts_str=", ".join(hosts),
+                                               keys_str=", ".join([
+                                                   "{value}->{state}".format(
+                                                       value=key.get('value'),
+                                                       state=str(
+                                                           key.get('state')))
+                                                   for key in keys]))
+                 )
+    return_value = True
+    with tempfile.NamedTemporaryFile(delete=True) as tmp_manifest:
+        # Generate manifest to be used for configuring the hosts
+        lines_list = [
+            '  --key "{key}" --state {state} \\\n'.format(
+                key=key['value'],
+                state='present' if key['state'] else 'absent'
+            ).encode('utf-8')
+            for key in keys]
+        lines_list.insert(0, b'__ssh_authorized_keys root \\\n')
+        tmp_manifest.writelines(lines_list)
+        tmp_manifest.flush()
+        try:
+            configure_hosts_simple(hosts,
+                                   tmp_manifest.name,
+                                   verbose=cdist.argparse.VERBOSE_TRACE)
+        except Exception as cdist_exception:
+            logger.error(cdist_exception)
+            return_value = False
+            email_data = {
+                'subject': "celery save_ssh_key error - task id {0}".format(
+                    self.request.id.__str__()),
+                'from_email': current_task.request.hostname,
+                'to': settings.DCL_ERROR_EMAILS_TO_LIST,
+                'body': "Task Id: {0}\nResult: {1}\nTraceback: {2}".format(
+                    self.request.id.__str__(), False, str(cdist_exception)),
+            }
+            send_plain_email_task(email_data)
+    return return_value
+
+
+@app.task
+def save_ssh_key_error_handler(uuid):
+    result = AsyncResult(uuid)
+    exc = result.get(propagate=False)
+    logger.error('Task {0} raised exception: {1!r}\n{2!r}'.format(
+        uuid, exc, result.traceback))
+    email_data = {
+        'subject': "[celery error] Save SSH key error {0}".format(uuid),
+        'from_email': current_task.request.hostname,
+        'to': settings.DCL_ERROR_EMAILS_TO_LIST,
+        'body': "Task Id: {0}\nResult: {1}\nTraceback: {2}".format(
+            uuid, exc, result.traceback),
+    }
+    send_plain_email_task(email_data)

utils/tests.py

@@ -1,16 +1,20 @@
 import uuid
+from time import sleep
 from unittest.mock import patch
 
 import stripe
+from celery.result import AsyncResult
+from django.conf import settings
 from django.http.request import HttpRequest
 from django.test import Client
-from django.test import TestCase
+from django.test import TestCase, override_settings
+from unittest import skipIf
 from model_mommy import mommy
 
 from datacenterlight.models import StripePlan
 from membership.models import StripeCustomer
 from utils.stripe_utils import StripeUtils
-from django.conf import settings
+from .tasks import save_ssh_key
 
 
 class BaseTestCase(TestCase):
@@ -235,3 +239,57 @@ class StripePlanTestCase(TestStripeCustomerDescription):
                 'response_object').stripe_plan_id}])
         self.assertIsNone(result.get('response_object'), None)
         self.assertIsNotNone(result.get('error'))
+
+
+class SaveSSHKeyTestCase(TestCase):
+    """
+    A test case to test the celery save_ssh_key task
+    """
+
+    @override_settings(
+        task_eager_propagates=True,
+        task_always_eager=True,
+    )
+    def setUp(self):
+        self.public_key = settings.TEST_MANAGE_SSH_KEY_PUBKEY
+        self.hosts = settings.TEST_MANAGE_SSH_KEY_HOST
+
+    @skipIf(settings.TEST_MANAGE_SSH_KEY_PUBKEY is None or
+            settings.TEST_MANAGE_SSH_KEY_PUBKEY == "" or
+            settings.TEST_MANAGE_SSH_KEY_HOST is None or
+            settings.TEST_MANAGE_SSH_KEY_HOST is "",
+            """Skipping test_save_ssh_key_add because either host
+             or public key were not specified or were empty""")
+    def test_save_ssh_key_add(self):
+        async_task = save_ssh_key.delay([self.hosts],
+                                        [{'value': self.public_key,
+                                          'state': True}])
+        save_ssh_key_result = None
+        for i in range(0, 10):
+            sleep(5)
+            res = AsyncResult(async_task.task_id)
+            if type(res.result) is bool:
+                save_ssh_key_result = res.result
+                break
+        self.assertIsNotNone(save_ssh_key, "save_ssh_key_result is None")
+        self.assertTrue(save_ssh_key_result, "save_ssh_key_result is False")
+
+    @skipIf(settings.TEST_MANAGE_SSH_KEY_PUBKEY is None or
+            settings.TEST_MANAGE_SSH_KEY_PUBKEY == "" or
+            settings.TEST_MANAGE_SSH_KEY_HOST is None or
+            settings.TEST_MANAGE_SSH_KEY_HOST is "",
+            """Skipping test_save_ssh_key_add because either host
+             or public key were not specified or were empty""")
+    def test_save_ssh_key_remove(self):
+        async_task = save_ssh_key.delay([self.hosts],
+                                        [{'value': self.public_key,
+                                          'state': False}])
+        save_ssh_key_result = None
+        for i in range(0, 10):
+            sleep(5)
+            res = AsyncResult(async_task.task_id)
+            if type(res.result) is bool:
+                save_ssh_key_result = res.result
+                break
+        self.assertIsNotNone(save_ssh_key, "save_ssh_key_result is None")
+        self.assertTrue(save_ssh_key_result, "save_ssh_key_result is False")