jumpserver/apps/terminal/automations/deploy_applet_host/__init__.py

import os
import datetime
import shutil

import yaml
from django.utils import timezone
from django.conf import settings

from common.utils import get_logger
from common.db.utils import safe_db_connection
from ops.ansible import PlaybookRunner, JMSInventory

logger = get_logger(__name__)
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))


class DeployAppletHostManager:
    def __init__(self, deployment):
        self.deployment = deployment
        self.run_dir = self.get_run_dir()

    @staticmethod
    def get_run_dir():
        base = os.path.join(settings.ANSIBLE_DIR, 'applet_host_deploy')
        now = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        return os.path.join(base, now)

    def generate_playbook(self):
        playbook_src = os.path.join(CURRENT_DIR, 'playbook.yml')
        with open(playbook_src) as f:
            plays = yaml.safe_load(f)
        for play in plays:
            play['vars'].update(self.deployment.host.deploy_options)
            play['vars']['DownloadHost'] = settings.BASE_URL + '/download/'
            play['vars']['CORE_HOST'] = settings.BASE_URL
            play['vars']['BOOTSTRAP_TOKEN'] = settings.BOOSTRAP_TOKEN
            play['vars']['HOST_NAME'] = self.deployment.host.name

        playbook_dir = os.path.join(self.run_dir, 'playbook')
        playbook_dst = os.path.join(playbook_dir, 'main.yml')
        os.makedirs(playbook_dir, exist_ok=True)
        with open(playbook_dst, 'w') as f:
            yaml.safe_dump(plays, f)
        return playbook_dst

    def generate_inventory(self):
        inventory = JMSInventory([self.deployment.host], account_policy='privileged_only')
        inventory_dir = os.path.join(self.run_dir, 'inventory')
        inventory_path = os.path.join(inventory_dir, 'hosts.yml')
        inventory.write_to_file(inventory_path)
        return inventory_path

    def _run(self, **kwargs):
        inventory = self.generate_inventory()
        playbook = self.generate_playbook()
        runner = PlaybookRunner(
            inventory=inventory, playbook=playbook, project_dir=self.run_dir
        )
        return runner.run(**kwargs)

    def run(self, **kwargs):
        try:
            self.deployment.date_start = timezone.now()
            cb = self._run(**kwargs)
            self.deployment.status = cb.status
        except Exception as e:
            logger.error("Error: {}".format(e))
            self.deployment.status = 'error'
        finally:
            self.deployment.date_finished = timezone.now()
            with safe_db_connection():
                self.deployment.save()