mirror of https://github.com/jumpserver/jumpserver
121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
import uuid
|
|
from celery import current_task
|
|
from django.db import models
|
|
from django.utils.translation import ugettext_lazy as _
|
|
|
|
from common.const.choices import Trigger
|
|
from common.mixins.models import CommonModelMixin
|
|
from common.db.fields import EncryptJsonDictTextField
|
|
from orgs.mixins.models import OrgModelMixin
|
|
from ops.mixin import PeriodTaskModelMixin
|
|
from assets.models import Node, Asset
|
|
from assets.tasks import execute_automation
|
|
from assets.const import AutomationTypes
|
|
|
|
|
|
class BaseAutomation(CommonModelMixin, PeriodTaskModelMixin, OrgModelMixin):
|
|
accounts = models.JSONField(default=list, verbose_name=_("Accounts"))
|
|
nodes = models.ManyToManyField('assets.Node', blank=True, verbose_name=_("Nodes"))
|
|
assets = models.ManyToManyField('assets.Asset', blank=True, verbose_name=_("Assets"))
|
|
type = models.CharField(max_length=16, choices=AutomationTypes.choices, verbose_name=_('Type'))
|
|
is_active = models.BooleanField(default=True, verbose_name=_("Is active"))
|
|
comment = models.TextField(blank=True, verbose_name=_('Comment'))
|
|
|
|
def __str__(self):
|
|
return self.name + '@' + str(self.created_by)
|
|
|
|
@classmethod
|
|
def generate_unique_name(cls, name):
|
|
while True:
|
|
name = name + str(uuid.uuid4())[:8]
|
|
try:
|
|
cls.objects.get(name=name)
|
|
except cls.DoesNotExist:
|
|
return name
|
|
|
|
def get_all_assets(self):
|
|
nodes = self.nodes.all()
|
|
node_asset_ids = Node.get_nodes_all_assets(*nodes).values_list('id', flat=True)
|
|
direct_asset_ids = self.assets.all().values_list('id', flat=True)
|
|
asset_ids = set(list(direct_asset_ids) + list(node_asset_ids))
|
|
return Asset.objects.filter(id__in=asset_ids)
|
|
|
|
def all_assets_group_by_platform(self):
|
|
assets = self.get_all_assets().prefetch_related('platform')
|
|
return assets.group_by_platform()
|
|
|
|
def get_register_task(self):
|
|
name = f"automation_{self.type}_strategy_period_{str(self.id)[:8]}"
|
|
task = execute_automation.name
|
|
args = (str(self.id), Trigger.timing, self.type)
|
|
kwargs = {}
|
|
return name, task, args, kwargs
|
|
|
|
def get_many_to_many_ids(self, field: str):
|
|
return [str(i) for i in getattr(self, field).all().values_list('id', flat=True)]
|
|
|
|
def to_attr_json(self):
|
|
return {
|
|
'name': self.name,
|
|
'type': self.type,
|
|
'org_id': str(self.org_id),
|
|
'comment': self.comment,
|
|
'accounts': self.accounts,
|
|
'nodes': self.get_many_to_many_ids('nodes'),
|
|
'assets': self.get_many_to_many_ids('assets'),
|
|
}
|
|
|
|
def execute(self, trigger=Trigger.manual):
|
|
try:
|
|
eid = current_task.request.id
|
|
except AttributeError:
|
|
eid = str(uuid.uuid4())
|
|
|
|
execution = self.executions.model.objects.create(
|
|
id=eid, trigger=trigger, automation=self,
|
|
snapshot=self.to_attr_json(),
|
|
)
|
|
return execution.start()
|
|
|
|
class Meta:
|
|
unique_together = [('org_id', 'name')]
|
|
verbose_name = _("Automation task")
|
|
|
|
|
|
class AutomationExecution(OrgModelMixin):
|
|
id = models.UUIDField(default=uuid.uuid4, primary_key=True)
|
|
automation = models.ForeignKey(
|
|
'BaseAutomation', related_name='executions', on_delete=models.CASCADE,
|
|
verbose_name=_('Automation task')
|
|
)
|
|
status = models.CharField(max_length=16, default='pending', verbose_name=_('Status'))
|
|
date_created = models.DateTimeField(auto_now_add=True, verbose_name=_('Date created'))
|
|
date_start = models.DateTimeField(null=True, verbose_name=_('Date start'), db_index=True)
|
|
date_finished = models.DateTimeField(null=True, verbose_name=_("Date finished"))
|
|
snapshot = EncryptJsonDictTextField(
|
|
default=dict, blank=True, null=True, verbose_name=_('Automation snapshot')
|
|
)
|
|
trigger = models.CharField(
|
|
max_length=128, default=Trigger.manual, choices=Trigger.choices,
|
|
verbose_name=_('Trigger mode')
|
|
)
|
|
|
|
class Meta:
|
|
verbose_name = _('Automation task execution')
|
|
|
|
@property
|
|
def manager_type(self):
|
|
return self.snapshot['type']
|
|
|
|
@property
|
|
def recipients(self):
|
|
recipients = self.snapshot.get('recipients')
|
|
if not recipients:
|
|
return []
|
|
return recipients.values()
|
|
|
|
def start(self):
|
|
from assets.automations.endpoint import ExecutionManager
|
|
manager = ExecutionManager(execution=self)
|
|
return manager.run()
|