from awx.main import utils
from awx.conf.migrations._reencrypt import (
    decrypt_field,
    should_decrypt_field,
)


__all__ = ['replace_aesecb_fernet']


def replace_aesecb_fernet(apps, schema_editor):
    _notification_templates(apps)
    _credentials(apps)
    _unified_jobs(apps)


def _notification_templates(apps):
    NotificationTemplate = apps.get_model('main', 'NotificationTemplate')
    for nt in NotificationTemplate.objects.all():
        for field in filter(lambda x: nt.notification_class.init_parameters[x]['type'] == "password",
                            nt.notification_class.init_parameters):
            if should_decrypt_field(nt.notification_configuration[field]):
                value = decrypt_field(nt, 'notification_configuration', subfield=field)
                nt.notification_configuration[field] = value
        nt.save()


def _credentials(apps):
    # this monkey-patch is necessary to make the implicit role generation save
    # signal use the correct Role model (the version active at this point in
    # migration, not the one at HEAD)
    orig_current_apps = utils.get_current_apps
    try:
        utils.get_current_apps = lambda: apps
        for credential in apps.get_model('main', 'Credential').objects.all():
            for field_name, value in credential.inputs.items():
                if should_decrypt_field(value):
                    value = decrypt_field(credential, field_name)
                    credential.inputs[field_name] = value
            credential.save()
    finally:
        utils.get_current_apps = orig_current_apps



def _unified_jobs(apps):
    UnifiedJob = apps.get_model('main', 'UnifiedJob')
    for uj in UnifiedJob.objects.all():
        if uj.start_args is not None:
            if should_decrypt_field(uj.start_args):
                start_args = decrypt_field(uj, 'start_args')
                uj.start_args = start_args
                uj.save()

@@ -1,16 +1,73 @@
+from django.utils.translation import ugettext_lazy as _%0A%0A
 from awx.main im
@@ -174,43 +174,989 @@
 %0A)%0A%0A
-%0A__all__ = %5B'replace_aesecb_fernet'
+from awx.main.notifications.email_backend import CustomEmailBackend%0Afrom awx.main.notifications.slack_backend import SlackBackend%0Afrom awx.main.notifications.twilio_backend import TwilioBackend%0Afrom awx.main.notifications.pagerduty_backend import PagerDutyBackend%0Afrom awx.main.notifications.hipchat_backend import HipChatBackend%0Afrom awx.main.notifications.webhook_backend import WebhookBackend%0Afrom awx.main.notifications.irc_backend import IrcBackend%0Afrom awx.main.models.credential import Credential%0A%0A__all__ = %5B'replace_aesecb_fernet'%5D%0A%0A%0ANOTIFICATION_TYPES = %5B('email', _('Email'), CustomEmailBackend),%0A                      ('slack', _('Slack'), SlackBackend),%0A                      ('twilio', _('Twilio'), TwilioBackend),%0A                      ('pagerduty', _('Pagerduty'), PagerDutyBackend),%0A                      ('hipchat', _('HipChat'), HipChatBackend),%0A                      ('webhook', _('Webhook'), WebhookBackend),%0A                      ('irc', _('IRC'), IrcBackend)
 %5D%0A%0A%0A
@@ -1437,32 +1437,198 @@
 .objects.all():%0A
+        CLASS_FOR_NOTIFICATION_TYPE = dict(%5B(x%5B0%5D, x%5B2%5D) for x in NOTIFICATION_TYPES%5D)%0A        notification_class = CLASS_FOR_NOTIFICATION_TYPE%5Bnt.notification_type%5D%0A
         for fiel
@@ -1649,19 +1649,16 @@
 mbda x: 
-nt.
 notifica
@@ -1730,35 +1730,32 @@
                 
-nt.
 notification_cla
@@ -2053,348 +2053,164 @@
   # 
-this monkey-patch is necessary to make the implicit role generation save%0A    # signal use the correct Role model (the version active at this point in%0A    # migration, not the one at HEAD)%0A    orig_current_apps = utils.get_current_apps%0A    try:%0A        utils.get_current_apps = lambda: apps%0A        for credential in apps.get_model('main', '
+TODO: Try to not use the model directly imported from our%0A    # source (should use apps.get_model) to make the migration less britle.%0A    for credential in 
 Cred
@@ -2215,18 +2215,16 @@
 edential
-')
 .objects
@@ -2235,28 +2235,24 @@
 ():%0A        
-   
- 
 for field_na
@@ -2295,36 +2295,32 @@
 ():%0A            
-    
 if should_decryp
@@ -2343,36 +2343,32 @@
                 
- 
-   
 value = decrypt_
@@ -2405,36 +2405,32 @@
                 
-   
- 
 credential.input
@@ -2451,36 +2451,32 @@
 = value%0A        
-    
 credential.save(
@@ -2481,73 +2481,8 @@
 e()%0A
-    finally:%0A        utils.get_current_apps = orig_current_apps%0A%0A
 %0A%0Ade

copy NotificationTemplate code for migration and use Credential model directly