<commit_before># FireStic - Python script for indexing FireEye json alerts
# into Elasticsearch over http...and some alerting too
#
# Please see: https://github.com/spcampbell/firestic
#
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from datetime import datetime
import pytz  		# pip install pytz
from premailer import transform  # pip install premailer
import pystache

def sendAlert(theJson, fsconfig):

    # Prepare alert data to include in email
    emailData = gatherEmailData(theJson, fsconfig.myTimezone)

    # Build html version of message
    htmlEmail = buildHTMLMessage(emailData)

    # Build text version of message
    textEmail = buildTextMessage(emailData)

    # Email subject
    subjectLine = "FireStic Alert - " + emailData['alertname'] + " - "
    subjectLine += emailData['alertid'] + " - " + emailData['action']

    # ---------------------------------------
    # !!!!!!!!! Issues creating SMS message - working on fix !!!!!!!!!
    # Send SMS
    #if emailData['alertname'] in fsconfig.smsTypeAlertOn:
    #    if emailData['action'] in fsconfig.smsActionAlertOn:
    #        txtMessages = splitForSMS(textEmail, emailData['alertid'])
    #        for txtMessage in txtMessages:
    #            msg = MIMEMultipart('alternative')
    #            msg['From'] = fsconfig.fromEmail
    #            msg['To'] = fsconfig.toSMS
    #            msg.attach(MIMEText(txtMessage, 'plain'))
    #            s = smtplib.SMTP(fsconfig.smtpServer, fsconfig.smtpPort)
    #            s.sendmail(fsconfig.fromEmail, fsconfig.toSMS, msg.as_string())
    #            s.quit()

    # ---------------------------------------

    # Send email
    if emailData['alertname'] in fsconfig.emailTypeAlertOn:
        msg = MIMEMultipart('alternative')
        msg['Subject'] = subjectLine
        msg['From'] = fsconfig.fromEmail
        msg['To'] = fsconfig.toEmail
        part1 = MIMEText(textEmail, 'plain')
        part2 = MIMEText(htmlEmail, 'html')
        msg.attach(part1)
        msg.attach(part2)
        s = smtplib.SMTP(fsconfig.smtpServer, fsconfig.smtpPort)
        s.sendmail(fsconfig.fromEmail, fsconfig.toEmail, msg.as_string())
        s.quit()

    # ---------------------------------------

def buildHTMLMessage(emailData):
    # load HTML email template from file
    tf = open('mustache_templates/template.html', 'r')
    htmlTemplate = tf.read()
    tf.close()

    # render html message from mustache template
    htmlEmail = pystache.render(htmlTemplate, emailData)
    htmlEmail = encode_for_html(htmlEmail)
    # the following uses premailer to move css inline
    htmlEmail = transform(htmlEmail)

    return htmlEmail

def buildTextMessage(emailData):
    # load text message template from file
    tf = open('mustache_templates/template.txt', 'r')
    textTemplate = tf.read()
    tf.close()

    # render text message
    textEmail = pystache.render(textTemplate, emailData)

    return textEmail

def encode_for_html(unicode_data, encoding='ascii'):
    try:
        return unicode_data.encode(encoding, 'xmlcharrefreplace')
    except:
        print "encode_for_html failed for: "
        print unicode_data
        return " "

def splitForSMS(fullText, alertID):
    texts = []
    words = fullText.split(' ')
    curtext = ''
    for word in words:
        # for the first word, drop the space
        if len(curtext) == 0:
            curtext += word

        # check if there's enough space left in the current message
        elif len(curtext) <= 155 - (len(word) + 1):
            curtext += ' ' + word

        # not enough space. make a new message
        else:
            texts.append(curtext)
            curtext = '(' + alertID + '...) ' + word
    if curtext != '':
        texts.append(curtext)

    return texts

def gatherEmailData(alertData, myTimezone):
    # TODO: refactor now that using mustache for template.
    # TODO: use get()

    # The process of gathering our data from the json sent over by FireEye is
    # complicated by the various permutations of how that data is presented.
    # It is likely that some of the following checks are redundant or
    # unnecessary. However, without a definitive answer on what all of the
    # possibilities are, the code tries to cover the worst case scenario.

    severityLevels = {u'crit': u'Critical', u'majr': u'Major', u'minr': u'Minor'}
    severityColors = {u'crit': u'red', u'majr':u'orange', u'minr':u'yellow'}
    emptyValue = u'N/A'
    emailData = {}

    emailData['timestamp'] = aslocaltimestr(alertData['@timestamp'], myTimezone)

    alertData = alertData.setdefault('alert', {})

    emailData['alertname'] = alertData.setdefault('name', emptyValue)
    emailData['alertid'] = str(alertData['id'])
    emailData['action'] = alertData.setdefault('action', emptyValue)

    emailData['severity'] = alertData.setdefault('severity', emptyValue)
    if emailData['severity'] in severityLevels:
        emailData['severitycolor'] = severityColors[emailData['severity']]
        emailData['severity'] = severityLevels[emailData['severity']]

    emailData['alerturl'] = alertData.setdefault('alert-url', emptyValue)

    # ips-events provide a signature name and possible a CVE-ID
    # everything else tries to provide info on the malware
    if emailData['alertname'] == u'ips-event':
        emailData['threatname'] = alertData.setdefault('explanation', {}).setdefault('ips-detected', {}).setdefault('sig-name', emptyValue)
        emailData['threatinfo'] = alertData.setdefault('explanation', {}).setdefault('ips-detected', {}).setdefault('cve-id', emptyValue)
        if (emailData['threatinfo'] != emptyValue) and (emailData['threatinfo'] is not None) and (emailData['threatinfo'] != ''):
            emailData['threatinfo'] = u'' + emailData['threatinfo'] + u' on cvedetails.com'
        else:
            emailData['threatinfo'] = emptyValue
        emailData['action'] = alertData.setdefault('explanation', {}).setdefault('ips-detected', {}).setdefault('action-taken', emptyValue)
    else:
        # sometimes the malware field is an array, sometimes not
        mwNames = []
        mwInfo = []
        #mwURLs = []
        urllist = []
        if isinstance(alertData.setdefault('explanation', {}).setdefault('malware-detected', {}).setdefault('malware', {}), list):
            for element in alertData['explanation']['malware-detected']['malware']:
                thisURLrow = {}
                # if (element.has_key('name')) and (element['name'] is not None):
                if ('name' in element) and (element['name'] is not None):
                    mwNames.append(element['name'])
                    thisURLrow['name'] = element['name']
                # if (element.has_key('original')) and (element['original'] is not None):
                if ('original' in element) and (element['original'] is not None):
                    mwNames.append(element['original'])
                # if (element.has_key('stype')) and (element['stype'] is not None):
                if ('stype' in element) and (element['stype'] is not None):
                    mwInfo.append(element['stype'])
                if ('url' in element) and (element['url'] is not None):
                    #mwURLs.append(element['url'])
                    thisURLrow['url'] = element['url']
                if ('objurl' in element) and (element['objurl'] is not None):
                    #mwURLs.append(element['objurl'])
                    thisURLrow['url'] = element['objurl']
                urllist.append(thisURLrow)
        else:
                element = alertData['explanation']['malware-detected']['malware']
                thisURLrow = {}
                # if (element.has_key('name')) and (element['name'] is not None):
                if ('name' in element) and (element['name'] is not None):
                    mwNames.append(element['name'])
                    thisURLrow['name'] = element['name']
                # if (element.has_key('original')) and (element['original'] is not None):
                if ('original' in element) and (element['original'] is not None):
                    mwNames.append(element['original'])
                # if (element.has_key('stype')) and (element['stype'] is not None):
                if ('stype' in element) and (element['stype'] is not None):
                    mwInfo.append(element['stype'])
                if ('url' in element) and (element['url'] is not None):
                    #mwURLs.append(element['url'])
                    thisURLrow['url'] = element['url']
                if ('objurl' in element) and (element['objurl'] is not None):
                    #mwURLs.append(element['objurl'])
                    thisURLrow['url'] = element['objurl']

            #emailData['threatname'] = alertData['explanation']['malware-detected']['malware'].setdefault('name', emptyValue)
            #emailData['threatinfo'] = alertData['explanation']['malware-detected']['malware'].setdefault('stype', emptyValue)

        if len(mwNames):
            emailData['threatname'] = ', '.join(mwNames)
        else:
            emailData['threatname'] = emptyValue
        if len(mwInfo):
            emailData['threatinfo'] = ', '.join(mwInfo)
        else:
            emailData['threatinfo'] = emptyValue

        emailData['threatURLs'] = []
        if len(urllist):
            emailData['threatURLs']['urllist'] = urllist
        #for url in mwURLs:
        #    thisDict = {'url':url}
        #    emailData['threatURLs']['urllist'].append(thisDict)

    emailData['sourceip'] = alertData.setdefault('src', {}).setdefault('ip', emptyValue)
    emailData['destinationip'] = alertData.setdefault('dst', {}).setdefault('ip', emptyValue)

    # source detail
    if alertData['src'].setdefault('geoip', {}) is None:
        alertData['src']['geoip'] = {}
    if alertData['src']['geoip'].setdefault('city', emptyValue) is None:
        alertData['src']['geoip']['city'] = emptyValue
    emailData['sourcecity'] = alertData['src']['geoip']['city']
    if alertData['src']['geoip'].setdefault('region_code', emptyValue) is None:
        alertData['src']['geoip']['region_code'] = emptyValue
    emailData['sourceregion'] = alertData['src']['geoip']['region_code']
    if alertData['src']['geoip'].setdefault('country_name', emptyValue) is None:
        alertData['src']['geoip']['country_name'] = emptyValue
    emailData['sourcecountry'] = alertData['src']['geoip']['country_name']
    if alertData['src']['geoip'].setdefault('asn', emptyValue) is None:
        alertData['src']['geoip']['asn'] = emptyValue
    emailData['sourceasn'] = alertData['src']['geoip']['asn']
    if alertData['src']['geoip'].setdefault('hostname', emptyValue) is None:
        alertData['src']['geoip']['hostname'] = emptyValue
    emailData['sourcehostname'] = alertData['src']['geoip']['hostname']

    # destination detail
    if alertData['dst'].setdefault('geoip', {}) is None:
        alertData['dst']['geoip'] = {}
    if alertData['dst']['geoip'].setdefault('city', emptyValue) is None:
        alertData['dst']['geoip']['city'] = emptyValue
    emailData['destinationcity'] = alertData['dst']['geoip']['city']
    if alertData['dst']['geoip'].setdefault('region_code', emptyValue) is None:
        alertData['dst']['geoip']['region_code'] = emptyValue
    emailData['destinationregion'] = alertData['dst']['geoip']['region_code']
    if alertData['dst']['geoip'].setdefault('country_name', emptyValue) is None:
        alertData['dst']['geoip']['country_name'] = emptyValue
    emailData['destinationcountry'] = alertData['dst']['geoip']['country_name']
    if alertData['dst']['geoip'].setdefault('asn', emptyValue) is None:
        alertData['dst']['geoip']['asn'] = emptyValue
    emailData['destinationasn'] = alertData['dst']['geoip']['asn']
    if alertData['dst']['geoip'].setdefault('hostname', emptyValue) is None:
        alertData['dst']['geoip']['hostname'] = emptyValue
    emailData['destinationhostname'] = alertData['dst']['geoip']['hostname']

    # instance['osinfo'] --> name of OS
    # instance['application']['app-name'] --> targeted application
    # instance['malicious-alert'] --> list of malicious activity
    #   example of one element in list:
    #       'msg':'Process is registering a hook to monitor...'
    #       'classtype':'Keylogging-Activity'
    #       'display-msg':'High-level keyboard hook registered'

    #{{ #summaryinfo }}
    #   {{ osinfo }}
    #   {{ app-name }}
    #   {{ #malicious-alert }}
    #       {{ classtype }}
    #       {{ msg }}
    #       {{ display-msg }}
    #   {{ /malicious-alert }}
    #{{ /summaryinfo }}

    # basic information concerning malicious activity from os-changes
    if (len(alertData['explanation'].setdefault('summaryinfo',[]))):
        emailData['summaryinfo'] = alertData['explanation']['summaryinfo']

    return emailData

def aslocaltimestr(utc_str, myTimezone):
    utc_dt = datetime.strptime(utc_str, '%Y-%m-%dT%H:%M:%S.%fZ')
    local_tz = pytz.timezone(myTimezone)
    local_dt = utc_dt.replace(tzinfo=pytz.utc).astimezone(local_tz)
    return local_tz.normalize(local_dt).strftime('%m-%d-%Y %H:%M:%S %Z%z')
<commit_msg>template data update<commit_after># FireStic - Python script for indexing FireEye json alerts
# into Elasticsearch over http...and some alerting too
#
# Please see: https://github.com/spcampbell/firestic
#
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from datetime import datetime
import pytz  		# pip install pytz
from premailer import transform  # pip install premailer
import pystache

def sendAlert(theJson, fsconfig):

    # Prepare alert data to include in email
    emailData = gatherEmailData(theJson, fsconfig.myTimezone)

    # Build html version of message
    htmlEmail = buildHTMLMessage(emailData)

    # Build text version of message
    textEmail = buildTextMessage(emailData)

    # Email subject
    subjectLine = "FireStic Alert - " + emailData['alertname'] + " - "
    subjectLine += emailData['alertid'] + " - " + emailData['action']

    # ---------------------------------------
    # !!!!!!!!! Issues creating SMS message - working on fix !!!!!!!!!
    # Send SMS
    #if emailData['alertname'] in fsconfig.smsTypeAlertOn:
    #    if emailData['action'] in fsconfig.smsActionAlertOn:
    #        txtMessages = splitForSMS(textEmail, emailData['alertid'])
    #        for txtMessage in txtMessages:
    #            msg = MIMEMultipart('alternative')
    #            msg['From'] = fsconfig.fromEmail
    #            msg['To'] = fsconfig.toSMS
    #            msg.attach(MIMEText(txtMessage, 'plain'))
    #            s = smtplib.SMTP(fsconfig.smtpServer, fsconfig.smtpPort)
    #            s.sendmail(fsconfig.fromEmail, fsconfig.toSMS, msg.as_string())
    #            s.quit()

    # ---------------------------------------

    # Send email
    if emailData['alertname'] in fsconfig.emailTypeAlertOn:
        msg = MIMEMultipart('alternative')
        msg['Subject'] = subjectLine
        msg['From'] = fsconfig.fromEmail
        msg['To'] = fsconfig.toEmail
        part1 = MIMEText(textEmail, 'plain')
        part2 = MIMEText(htmlEmail, 'html')
        msg.attach(part1)
        msg.attach(part2)
        s = smtplib.SMTP(fsconfig.smtpServer, fsconfig.smtpPort)
        s.sendmail(fsconfig.fromEmail, fsconfig.toEmail, msg.as_string())
        s.quit()

    # ---------------------------------------

def buildHTMLMessage(emailData):
    # load HTML email template from file
    tf = open('mustache_templates/template.html', 'r')
    htmlTemplate = tf.read()
    tf.close()

    # render html message from mustache template
    htmlEmail = pystache.render(htmlTemplate, emailData)
    htmlEmail = encode_for_html(htmlEmail)
    # the following uses premailer to move css inline
    htmlEmail = transform(htmlEmail)

    return htmlEmail

def buildTextMessage(emailData):
    # load text message template from file
    tf = open('mustache_templates/template.txt', 'r')
    textTemplate = tf.read()
    tf.close()

    # render text message
    textEmail = pystache.render(textTemplate, emailData)

    return textEmail

def encode_for_html(unicode_data, encoding='ascii'):
    try:
        return unicode_data.encode(encoding, 'xmlcharrefreplace')
    except:
        print "encode_for_html failed for: "
        print unicode_data
        return " "

def splitForSMS(fullText, alertID):
    texts = []
    words = fullText.split(' ')
    curtext = ''
    for word in words:
        # for the first word, drop the space
        if len(curtext) == 0:
            curtext += word

        # check if there's enough space left in the current message
        elif len(curtext) <= 155 - (len(word) + 1):
            curtext += ' ' + word

        # not enough space. make a new message
        else:
            texts.append(curtext)
            curtext = '(' + alertID + '...) ' + word
    if curtext != '':
        texts.append(curtext)

    return texts

def gatherEmailData(alertData, myTimezone):
    # TODO: refactor now that using mustache for template.
    # TODO: use get()

    # The process of gathering our data from the json sent over by FireEye is
    # complicated by the various permutations of how that data is presented.
    # It is likely that some of the following checks are redundant or
    # unnecessary. However, without a definitive answer on what all of the
    # possibilities are, the code tries to cover the worst case scenario.

    severityLevels = {u'crit': u'Critical', u'majr': u'Major', u'minr': u'Minor'}
    severityColors = {u'crit': u'red', u'majr':u'orange', u'minr':u'yellow'}
    emptyValue = u'N/A'
    emailData = {}

    emailData['timestamp'] = aslocaltimestr(alertData['@timestamp'], myTimezone)

    alertData = alertData.setdefault('alert', {})

    emailData['alertname'] = alertData.setdefault('name', emptyValue)
    emailData['alertid'] = str(alertData['id'])
    emailData['action'] = alertData.setdefault('action', emptyValue)

    emailData['severity'] = alertData.setdefault('severity', emptyValue)
    if emailData['severity'] in severityLevels:
        emailData['severitycolor'] = severityColors[emailData['severity']]
        emailData['severity'] = severityLevels[emailData['severity']]

    emailData['alerturl'] = alertData.setdefault('alert-url', emptyValue)

    # ips-events provide a signature name and possible a CVE-ID
    # everything else tries to provide info on the malware
    if emailData['alertname'] == u'ips-event':
        emailData['threatname'] = alertData.setdefault('explanation', {}).setdefault('ips-detected', {}).setdefault('sig-name', emptyValue)
        emailData['threatinfo'] = alertData.setdefault('explanation', {}).setdefault('ips-detected', {}).setdefault('cve-id', emptyValue)
        if (emailData['threatinfo'] != emptyValue) and (emailData['threatinfo'] is not None) and (emailData['threatinfo'] != ''):
            emailData['threatinfo'] = u'' + emailData['threatinfo'] + u' on cvedetails.com'
        else:
            emailData['threatinfo'] = emptyValue
        emailData['action'] = alertData.setdefault('explanation', {}).setdefault('ips-detected', {}).setdefault('action-taken', emptyValue)
    else:
        # sometimes the malware field is an array, sometimes not
        mwNames = []
        mwInfo = []
        #mwURLs = []
        urllist = []
        if isinstance(alertData.setdefault('explanation', {}).setdefault('malware-detected', {}).setdefault('malware', {}), list):
            for element in alertData['explanation']['malware-detected']['malware']:
                thisURLrow = {}
                # if (element.has_key('name')) and (element['name'] is not None):
                if ('name' in element) and (element['name'] is not None):
                    mwNames.append(element['name'])
                    thisURLrow['name'] = element['name']
                # if (element.has_key('original')) and (element['original'] is not None):
                if ('original' in element) and (element['original'] is not None):
                    mwNames.append(element['original'])
                # if (element.has_key('stype')) and (element['stype'] is not None):
                if ('stype' in element) and (element['stype'] is not None):
                    mwInfo.append(element['stype'])
                if ('url' in element) and (element['url'] is not None):
                    #mwURLs.append(element['url'])
                    thisURLrow['url'] = element['url']
                if ('objurl' in element) and (element['objurl'] is not None):
                    #mwURLs.append(element['objurl'])
                    thisURLrow['url'] = element['objurl']
                urllist.append(thisURLrow)
        else:
                element = alertData['explanation']['malware-detected']['malware']
                thisURLrow = {}
                # if (element.has_key('name')) and (element['name'] is not None):
                if ('name' in element) and (element['name'] is not None):
                    mwNames.append(element['name'])
                    thisURLrow['name'] = element['name']
                # if (element.has_key('original')) and (element['original'] is not None):
                if ('original' in element) and (element['original'] is not None):
                    mwNames.append(element['original'])
                # if (element.has_key('stype')) and (element['stype'] is not None):
                if ('stype' in element) and (element['stype'] is not None):
                    mwInfo.append(element['stype'])
                if ('url' in element) and (element['url'] is not None):
                    #mwURLs.append(element['url'])
                    thisURLrow['url'] = element['url']
                if ('objurl' in element) and (element['objurl'] is not None):
                    #mwURLs.append(element['objurl'])
                    thisURLrow['url'] = element['objurl']

            #emailData['threatname'] = alertData['explanation']['malware-detected']['malware'].setdefault('name', emptyValue)
            #emailData['threatinfo'] = alertData['explanation']['malware-detected']['malware'].setdefault('stype', emptyValue)

        if len(mwNames):
            emailData['threatname'] = ', '.join(mwNames)
        else:
            emailData['threatname'] = emptyValue
        if len(mwInfo):
            emailData['threatinfo'] = ', '.join(mwInfo)
        else:
            emailData['threatinfo'] = emptyValue

        emailData['threatURLs'] = []
        if len(urllist):
            emailData['threatURLs'][] = {'urllist':urllist}
        #for url in mwURLs:
        #    thisDict = {'url':url}
        #    emailData['threatURLs']['urllist'].append(thisDict)

    emailData['sourceip'] = alertData.setdefault('src', {}).setdefault('ip', emptyValue)
    emailData['destinationip'] = alertData.setdefault('dst', {}).setdefault('ip', emptyValue)

    # source detail
    if alertData['src'].setdefault('geoip', {}) is None:
        alertData['src']['geoip'] = {}
    if alertData['src']['geoip'].setdefault('city', emptyValue) is None:
        alertData['src']['geoip']['city'] = emptyValue
    emailData['sourcecity'] = alertData['src']['geoip']['city']
    if alertData['src']['geoip'].setdefault('region_code', emptyValue) is None:
        alertData['src']['geoip']['region_code'] = emptyValue
    emailData['sourceregion'] = alertData['src']['geoip']['region_code']
    if alertData['src']['geoip'].setdefault('country_name', emptyValue) is None:
        alertData['src']['geoip']['country_name'] = emptyValue
    emailData['sourcecountry'] = alertData['src']['geoip']['country_name']
    if alertData['src']['geoip'].setdefault('asn', emptyValue) is None:
        alertData['src']['geoip']['asn'] = emptyValue
    emailData['sourceasn'] = alertData['src']['geoip']['asn']
    if alertData['src']['geoip'].setdefault('hostname', emptyValue) is None:
        alertData['src']['geoip']['hostname'] = emptyValue
    emailData['sourcehostname'] = alertData['src']['geoip']['hostname']

    # destination detail
    if alertData['dst'].setdefault('geoip', {}) is None:
        alertData['dst']['geoip'] = {}
    if alertData['dst']['geoip'].setdefault('city', emptyValue) is None:
        alertData['dst']['geoip']['city'] = emptyValue
    emailData['destinationcity'] = alertData['dst']['geoip']['city']
    if alertData['dst']['geoip'].setdefault('region_code', emptyValue) is None:
        alertData['dst']['geoip']['region_code'] = emptyValue
    emailData['destinationregion'] = alertData['dst']['geoip']['region_code']
    if alertData['dst']['geoip'].setdefault('country_name', emptyValue) is None:
        alertData['dst']['geoip']['country_name'] = emptyValue
    emailData['destinationcountry'] = alertData['dst']['geoip']['country_name']
    if alertData['dst']['geoip'].setdefault('asn', emptyValue) is None:
        alertData['dst']['geoip']['asn'] = emptyValue
    emailData['destinationasn'] = alertData['dst']['geoip']['asn']
    if alertData['dst']['geoip'].setdefault('hostname', emptyValue) is None:
        alertData['dst']['geoip']['hostname'] = emptyValue
    emailData['destinationhostname'] = alertData['dst']['geoip']['hostname']

    # instance['osinfo'] --> name of OS
    # instance['application']['app-name'] --> targeted application
    # instance['malicious-alert'] --> list of malicious activity
    #   example of one element in list:
    #       'msg':'Process is registering a hook to monitor...'
    #       'classtype':'Keylogging-Activity'
    #       'display-msg':'High-level keyboard hook registered'

    #{{ #summaryinfo }}
    #   {{ osinfo }}
    #   {{ app-name }}
    #   {{ #malicious-alert }}
    #       {{ classtype }}
    #       {{ msg }}
    #       {{ display-msg }}
    #   {{ /malicious-alert }}
    #{{ /summaryinfo }}

    # basic information concerning malicious activity from os-changes
    if (len(alertData['explanation'].setdefault('summaryinfo',[]))):
        emailData['summaryinfo'] = alertData['explanation']['summaryinfo']

    return emailData

def aslocaltimestr(utc_str, myTimezone):
    utc_dt = datetime.strptime(utc_str, '%Y-%m-%dT%H:%M:%S.%fZ')
    local_tz = pytz.timezone(myTimezone)
    local_dt = utc_dt.replace(tzinfo=pytz.utc).astimezone(local_tz)
    return local_tz.normalize(local_dt).strftime('%m-%d-%Y %H:%M:%S %Z%z')
<|endoftext|>


<commit_before>import discord
from discord.ext import commands
from cogs.utils.helpers import *
from cogs.utils.clip import *
from __main__ import settings, botdata
from cogs.utils import checks
import asyncio
import os
import string
import queue
import random
import re
import urllib.request
from random import randint
from .mangocog import *
from ctypes.util import find_library

discord.opus.load_opus(find_library('opus'))

def get_clipdirs():
	result = []
	for root, dirs, files in os.walk(settings.resourcedir + "clips/"):
		for d in dirs:
			result.append(d)
	result.sort()
	return result

# gets a list of all the mp3s in the indicated clipdir
def get_playlist(clipdir):
	clips = []
	for root, dirs, files in os.walk(settings.resourcedir + "clips/" + clipdir):
		for file in files:
			if file.endswith(".mp3"):
				clips.append(file[:-4])
	clips.sort()
	return clips

def remove_if_temp(mp3name):
	if os.path.isfile(mp3name):
		if os.path.dirname(mp3name) == os.path.join(settings.resourcedir, "temp"):
			os.remove(mp3name)
			print("removed temp file " + mp3name)

class Audio(MangoCog):
	"""Commands used to play audio
	"""

	def __init__(self, bot):
		MangoCog.__init__(self, bot)
		self.voice = None
		self.player = None
		self.clipqueue = queue.Queue()
		self.last_clip = None

	# whether or not the bot is currently talking
	def is_talking(self):
		return (self.player is not None) and (not self.player.is_done())

	def done_talking(self):
		self.player = None
		if not self.clipqueue.empty():
			self.play_next_clip()

	# gets the next clip from the clip queue
	def next_clip(self):
		if not self.clipqueue.empty():
			return self.clipqueue.get()
		raise ValueError("clip queue was empty when we tried to get the next one")

	# plays the next clip in the queue
	def play_next_clip(self):
		try:
			clip = self.next_clip()
			self.player = self.voice.create_ffmpeg_player(clip.audiopath, after=self.done_talking)
			self.player.volume = clip.volume
			self.player.start()
			print("playing: " + clip.audiopath)
			if self.last_clip != None and clip.audiopath != self.last_clip.audiopath:
				remove_if_temp(self.last_clip.audiopath)
			self.last_clip = clip
		except Exception as e:
			print(str(e))

	# try queueing an mp3 to play
	async def queue_clip(self, clip):
		if(self.voice is None):
			print("tried to talk while not in voice channel")
			await self.bot.say("not in voice channel m8")
			return

		self.clipqueue.put(clip)

		if not self.is_talking():
			self.play_next_clip()

	@commands.command(pass_context=True)
	async def play(self, ctx, clip : str=""):
		"""Plays an audio clip

		example:
		?play hello

		for a complete list of the available clips, try ?playlist"""
		try:
			await self.play_clip("local:" + clip)
		except ClipNotFound:
			await self.bot.say("'" + clip + "' is not a valid clip. 🤦 Try ?playlist. ")

	@commands.command(pass_context=True)
	async def playlist(self, ctx, section : str=None):
		"""Lists the audio clips available for ?play

		Calling this command with no arguments gets you a list of sections and a list of all of the clips

		To get the clips in a specific section, do ?playlist <section>"""
		dirs = get_clipdirs()

		message = ""
		clips = []

		if section is None:
			message += "**Sections:**\n"
			for section in dirs:
				message += "`{}` ".format(section)
			message += "\n**Clips:**\n"
			for section in dirs:
				clips += get_playlist(section)
		elif section not in dirs:
			message +=("Dats not a valid section. You can choose from one of these:\n")
			for section in dirs:
				message += "`{}` ".format(section)
		else:
			clips = get_playlist(section)

		if len(clips) > 0:
			clips.sort()
			for clip in clips:
				message += "`{}` ".format(clip)

		await self.bot.say(message)

	@commands.command(pass_context=True)
	async def playurl(self, ctx, mp3url : str):
		"""Plays an mp3 file at a url

		Make sure to use http, not https"""
		await self.play_clip("url:" + mp3url)

	@commands.command(pass_context=True)
	async def stop(self, ctx):
		"""Stops the currently playing audio

		Also empties the clip queue
		"""
		while not self.clipqueue.empty():
			try:
				self.clipqueue.get()
			except Empty:
				continue
		if not self.player is None:
			self.player.stop();

	@commands.command(pass_context=True)
	async def replay(self, ctx):
		"""Replays the last played clip
		"""
		if self.last_clip == None:
			await self.bot.say("Nobody said anythin' yet")
			return

		# If its not a temp file
		await self.bot.say("Replaying " + self.last_clip.clipid)

		await self.queue_clip(self.last_clip)

	@commands.command(pass_context=True)
	async def clipinfo(self, ctx, clipid=None):
		"""Gets information and a file for the given clip

		Not giving a clipid will print info about the last clip played

		clipid is specified like this:
		local:shitpickle
		dota:timb_ally_01
		"""
		if clipid is None:
			if self.last_clip == None:
				await self.bot.say("Nobody said anythin' yet")
				return
			clipid = self.last_clip.clipid

		clip = await self.get_clip(clipid)

		await self.bot.send_typing(ctx.message.channel)

		filename = clip.name

		if re.search(r"[^a-zA-Z0-9_]", filename):
			filename = clip.type()

		filename += os.path.splitext(clip.audiopath)[1]

		try:
			await self.bot.send_file(ctx.message.channel, clip.audiopath, filename=filename, content=await clip.get_info())
		except FileNotFoundError as e:
			# The file is probably actually a url
			fp = urllib.request.urlopen(clip.audiopath)
			await self.bot.send_file(ctx.message.channel, fp, filename=filename, content=await clip.get_info())
			fp.close()

	@commands.command(pass_context=True)
	async def setintro(self, ctx, clipname : str, user: discord.User=None):
		"""Sets your intro clip

		The argument is the name of the clip that will introduce you, for example:
		?setintro math
		Note: your intro clip cannot be longer than 4 seconds
		"""
		if user is None:
			user = ctx.message.author
		else:
			if not await checks.is_owner_check(ctx):
				await self.bot.say("You aint the boss of me 😠")
				return

		clip = await self.get_clip_try_types(clipname, "local|dota")

		audiolength = clip.audiolength

		if audiolength > 3.0:
			await self.bot.say("Dat clip is " + str(audiolength) + " seconds long, and intros gotta be less than 3.")
			return

		botdata.userinfo(user.id).intro = clip.clipid
		await self.bot.say("Yer intro is now " + clip.clipid)

	@commands.command(pass_context=True)
	async def setoutro(self, ctx, clipname : str, user: discord.User=None):
		"""Sets your outro clip

		The argument is the name of the clip that will 'outroduce' you, for example:
		?setoutro math
		Note: your outro clip cannot be longer than 4 seconds
		"""
		if user is None:
			user = ctx.message.author
		else:
			if not await checks.is_owner_check(ctx):
				await self.bot.say("You aint the boss of me 😠")
				return

		clip = await self.get_clip_try_types(clipname, "local|dota")

		audiolength = clip.audiolength

		if audiolength > 3.0:
			await self.bot.say("Dat clip is " + str(audiolength) + " seconds long, and outros gotta be less than 3.")
			return

		botdata.userinfo(user.id).outro = clip.clipid
		await self.bot.say("Yer outro is now " + clip.clipid)

	@commands.command(pass_context=True)
	async def tts(self, ctx, *, message : str):
		"""Like ?echo but for people who can't read

		...what more could you possibly need to know...
		"""
		fixes_dict = { "im": "i'm",
			"shes": "she's",
			"hes": "he's",
			"theyre": "they're"}
		text = ctx.message.clean_content[5:]
		for key in fixes_dict:
			text = re.sub(key, fixes_dict[key], text, re.IGNORECASE)
		await self.play_clip("tts:" + text)

	@commands.command(pass_context=True)
	async def later(self, ctx):
		"""Tells you how much later it is

		Theres 19 different ones"""
		await self.play_clip("local:later{}".format(randint(1,19)))

	# fixes discord user names which either are in all caps or have a number serving as a letter
	async def fix_name(self, name):
		# If all upper case or numbers n stuff, make all lower case
		if re.match(r"^[^a-z]*$", name):
			name = name.lower()

		# Simple replacing of all numbers that appear to be acting as letters
		letternumbers = { "0": "o", "1": "i", "3": "e", "4": "a", "7": "t" }
		for num in letternumbers:
			name = re.sub("{}([A-Za-z])".format(num), r"{}\1".format(letternumbers[num]), name)
		return name

	#function called when this event occurs
	async def on_voice_state_update(self, before, after):
		if self.voice is None or before.voice_channel == after.voice_channel:
			return # if the member didnt change channels, dont worry about it
		if after.voice_channel == self.voice.channel:
			print(after.name + " joined the channel")

			text = await self.fix_name(after.name)
			introclip = "local:helloits"

			userinfo = botdata.userinfo(after.id)
			if userinfo.intro != "" and userinfo.intro != introclip:
				introclip = userinfo.intro
				text = "its " + after.name

			await asyncio.sleep(3)
			await self.play_clip(introclip)
			await self.play_clip("tts:" + text)
		if before.voice_channel == self.voice.channel:
			print(before.name + " left the channel")

			text = await self.fix_name(before.name) + " has left!"
			outroclip = "local:farewell"

			userinfo = botdata.userinfo(before.id)
			if userinfo.outro != "" and userinfo.outro != outroclip:
				outroclip = userinfo.outro

			await asyncio.sleep(0.5)
			await self.play_clip(outroclip)
			await self.play_clip("tts:" + text)

def setup(bot):
	bot.add_cog(Audio(bot))<commit_msg>calling setintro and setoutro without arguments now has mangobyte tell you what your intro/outro is currently set to<commit_after>import discord
from discord.ext import commands
from cogs.utils.helpers import *
from cogs.utils.clip import *
from __main__ import settings, botdata
from cogs.utils import checks
import asyncio
import os
import string
import queue
import random
import re
import urllib.request
from random import randint
from .mangocog import *
from ctypes.util import find_library

discord.opus.load_opus(find_library('opus'))

def get_clipdirs():
	result = []
	for root, dirs, files in os.walk(settings.resourcedir + "clips/"):
		for d in dirs:
			result.append(d)
	result.sort()
	return result

# gets a list of all the mp3s in the indicated clipdir
def get_playlist(clipdir):
	clips = []
	for root, dirs, files in os.walk(settings.resourcedir + "clips/" + clipdir):
		for file in files:
			if file.endswith(".mp3"):
				clips.append(file[:-4])
	clips.sort()
	return clips

def remove_if_temp(mp3name):
	if os.path.isfile(mp3name):
		if os.path.dirname(mp3name) == os.path.join(settings.resourcedir, "temp"):
			os.remove(mp3name)
			print("removed temp file " + mp3name)

class Audio(MangoCog):
	"""Commands used to play audio
	"""

	def __init__(self, bot):
		MangoCog.__init__(self, bot)
		self.voice = None
		self.player = None
		self.clipqueue = queue.Queue()
		self.last_clip = None

	# whether or not the bot is currently talking
	def is_talking(self):
		return (self.player is not None) and (not self.player.is_done())

	def done_talking(self):
		self.player = None
		if not self.clipqueue.empty():
			self.play_next_clip()

	# gets the next clip from the clip queue
	def next_clip(self):
		if not self.clipqueue.empty():
			return self.clipqueue.get()
		raise ValueError("clip queue was empty when we tried to get the next one")

	# plays the next clip in the queue
	def play_next_clip(self):
		try:
			clip = self.next_clip()
			self.player = self.voice.create_ffmpeg_player(clip.audiopath, after=self.done_talking)
			self.player.volume = clip.volume
			self.player.start()
			print("playing: " + clip.audiopath)
			if self.last_clip != None and clip.audiopath != self.last_clip.audiopath:
				remove_if_temp(self.last_clip.audiopath)
			self.last_clip = clip
		except Exception as e:
			print(str(e))

	# try queueing an mp3 to play
	async def queue_clip(self, clip):
		if(self.voice is None):
			print("tried to talk while not in voice channel")
			await self.bot.say("not in voice channel m8")
			return

		self.clipqueue.put(clip)

		if not self.is_talking():
			self.play_next_clip()

	@commands.command(pass_context=True)
	async def play(self, ctx, clip : str=""):
		"""Plays an audio clip

		example:
		?play hello

		for a complete list of the available clips, try ?playlist"""
		try:
			await self.play_clip("local:" + clip)
		except ClipNotFound:
			await self.bot.say("'" + clip + "' is not a valid clip. 🤦 Try ?playlist. ")

	@commands.command(pass_context=True)
	async def playlist(self, ctx, section : str=None):
		"""Lists the audio clips available for ?play

		Calling this command with no arguments gets you a list of sections and a list of all of the clips

		To get the clips in a specific section, do ?playlist <section>"""
		dirs = get_clipdirs()

		message = ""
		clips = []

		if section is None:
			message += "**Sections:**\n"
			for section in dirs:
				message += "`{}` ".format(section)
			message += "\n**Clips:**\n"
			for section in dirs:
				clips += get_playlist(section)
		elif section not in dirs:
			message +=("Dats not a valid section. You can choose from one of these:\n")
			for section in dirs:
				message += "`{}` ".format(section)
		else:
			clips = get_playlist(section)

		if len(clips) > 0:
			clips.sort()
			for clip in clips:
				message += "`{}` ".format(clip)

		await self.bot.say(message)

	@commands.command(pass_context=True)
	async def playurl(self, ctx, mp3url : str):
		"""Plays an mp3 file at a url

		Make sure to use http, not https"""
		await self.play_clip("url:" + mp3url)

	@commands.command(pass_context=True)
	async def stop(self, ctx):
		"""Stops the currently playing audio

		Also empties the clip queue
		"""
		while not self.clipqueue.empty():
			try:
				self.clipqueue.get()
			except Empty:
				continue
		if not self.player is None:
			self.player.stop();

	@commands.command(pass_context=True)
	async def replay(self, ctx):
		"""Replays the last played clip
		"""
		if self.last_clip == None:
			await self.bot.say("Nobody said anythin' yet")
			return

		# If its not a temp file
		await self.bot.say("Replaying " + self.last_clip.clipid)

		await self.queue_clip(self.last_clip)

	@commands.command(pass_context=True)
	async def clipinfo(self, ctx, clipid=None):
		"""Gets information and a file for the given clip

		Not giving a clipid will print info about the last clip played

		clipid is specified like this:
		local:shitpickle
		dota:timb_ally_01
		"""
		if clipid is None:
			if self.last_clip == None:
				await self.bot.say("Nobody said anythin' yet")
				return
			clipid = self.last_clip.clipid

		clip = await self.get_clip(clipid)

		await self.bot.send_typing(ctx.message.channel)

		filename = clip.name

		if re.search(r"[^a-zA-Z0-9_]", filename):
			filename = clip.type()

		filename += os.path.splitext(clip.audiopath)[1]

		try:
			await self.bot.send_file(ctx.message.channel, clip.audiopath, filename=filename, content=await clip.get_info())
		except FileNotFoundError as e:
			# The file is probably actually a url
			fp = urllib.request.urlopen(clip.audiopath)
			await self.bot.send_file(ctx.message.channel, fp, filename=filename, content=await clip.get_info())
			fp.close()

	@commands.command(pass_context=True)
	async def setintro(self, ctx, clipname : str=None, user: discord.User=None):
		"""Sets your intro clip

		Calling this command without a clipname will tell you your current intro

		The argument is the name of the clip that will introduce you, for example:
		?setintro math
		Note: your intro clip cannot be longer than 4 seconds
		"""
		if user is None:
			user = ctx.message.author
		else:
			if not await checks.is_owner_check(ctx):
				await self.bot.say("You aint the boss of me 😠")
				return

		if clipname is None:
			intro = botdata.userinfo(user.id).intro
			if intro is None or intro == "":
				await self.bot.say("Yer intro isn't set. Try doin somethin' like `?setintro dota:gyro_items_01`")
				return
			else:
				await self.bot.say("Your intro is: {}".format(intro))
				await self.play_clip("tts:your intro is")
				await self.play_clip(intro)
				return

		clip = await self.get_clip_try_types(clipname, "local|dota")

		audiolength = clip.audiolength

		if audiolength > 3.0:
			await self.bot.say("Dat clip is " + str(audiolength) + " seconds long, and intros gotta be less than 3.")
			return

		botdata.userinfo(user.id).intro = clip.clipid
		await self.bot.say("Yer intro is now " + clip.clipid)

	@commands.command(pass_context=True)
	async def setoutro(self, ctx, clipname : str=None, user: discord.User=None):
		"""Sets your outro clip

		The argument is the name of the clip that will 'outroduce' you, for example:
		?setoutro math
		Note: your outro clip cannot be longer than 4 seconds
		"""
		if user is None:
			user = ctx.message.author
		else:
			if not await checks.is_owner_check(ctx):
				await self.bot.say("You aint the boss of me 😠")
				return

		if clipname is None:
			outro = botdata.userinfo(user.id).outro
			if outro is None or outro == "":
				await self.bot.say("Yer outro isn't set. Try doin somethin' like `?setoutro dota:troll_lose_03`")
				return
			else:
				await self.bot.say("Your outro is: {}".format(outro))
				await self.play_clip("tts:your outro is")
				await self.play_clip(outro)
				return

		clip = await self.get_clip_try_types(clipname, "local|dota")

		audiolength = clip.audiolength

		if audiolength > 3.0:
			await self.bot.say("Dat clip is " + str(audiolength) + " seconds long, and outros gotta be less than 3.")
			return

		botdata.userinfo(user.id).outro = clip.clipid
		await self.bot.say("Yer outro is now " + clip.clipid)

	@commands.command(pass_context=True)
	async def tts(self, ctx, *, message : str):
		"""Like ?echo but for people who can't read

		...what more could you possibly need to know...
		"""
		fixes_dict = { "im": "i'm",
			"shes": "she's",
			"hes": "he's",
			"theyre": "they're"}
		text = ctx.message.clean_content[5:]
		for key in fixes_dict:
			text = re.sub(key, fixes_dict[key], text, re.IGNORECASE)
		await self.play_clip("tts:" + text)

	@commands.command(pass_context=True)
	async def later(self, ctx):
		"""Tells you how much later it is

		Theres 19 different ones"""
		await self.play_clip("local:later{}".format(randint(1,19)))

	# fixes discord user names which either are in all caps or have a number serving as a letter
	async def fix_name(self, name):
		# If all upper case or numbers n stuff, make all lower case
		if re.match(r"^[^a-z]*$", name):
			name = name.lower()

		# Simple replacing of all numbers that appear to be acting as letters
		letternumbers = { "0": "o", "1": "i", "3": "e", "4": "a", "7": "t" }
		for num in letternumbers:
			name = re.sub("{}([A-Za-z])".format(num), r"{}\1".format(letternumbers[num]), name)
		return name

	#function called when this event occurs
	async def on_voice_state_update(self, before, after):
		if self.voice is None or before.voice_channel == after.voice_channel:
			return # if the member didnt change channels, dont worry about it
		if after.voice_channel == self.voice.channel:
			print(after.name + " joined the channel")

			text = await self.fix_name(after.name)
			introclip = "local:helloits"

			userinfo = botdata.userinfo(after.id)
			if userinfo.intro != "" and userinfo.intro != introclip:
				introclip = userinfo.intro
				text = "its " + after.name

			await asyncio.sleep(3)
			await self.play_clip(introclip)
			await self.play_clip("tts:" + text)
		if before.voice_channel == self.voice.channel:
			print(before.name + " left the channel")

			text = await self.fix_name(before.name) + " has left!"
			outroclip = "local:farewell"

			userinfo = botdata.userinfo(before.id)
			if userinfo.outro != "" and userinfo.outro != outroclip:
				outroclip = userinfo.outro

			await asyncio.sleep(0.5)
			await self.play_clip(outroclip)
			await self.play_clip("tts:" + text)

def setup(bot):
	bot.add_cog(Audio(bot))<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
from datetime import datetime

from dateutil.relativedelta import relativedelta

from django.contrib import messages
from django.contrib.auth.decorators import login_required
from django.core.urlresolvers import reverse
from django.http import (
    Http404,
    HttpResponse,
)
from django.views.generic import (
    CreateView,
    DeleteView,
    DetailView,
    ListView,
    UpdateView,
)

from braces.views import (
    LoginRequiredMixin,
    StaffuserRequiredMixin,
)

from base.view_utils import BaseMixin

from .forms import (
    BookingForm,
    BookingNotesStaffForm,
    BookingNotesUserForm,
    CategoryForm,
    LocationForm,
    RotaEmptyForm,
    RotaForm,
    RotaTypeForm,
)
from .models import (
    Booking,
    BookingSettings,
    Category,
    Location,
    Rota,
    RotaType,
)
from .service import PdfCalendar

def _url_booking(booking):
    """If we are editing from the detail form."""
    if BookingSettings.load().edit_from_detail:
        return reverse('booking.detail', args=[booking.pk])
    else:
        return _url_booking_list_month(booking)

def _url_booking_list_month(booking):
    return reverse(
        'booking.list.month',
        kwargs=dict(
            year=booking.start_date.year,
            month=booking.start_date.month
        )
    )

class BookingCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    model = Booking

    def get_success_url(self):
        return _url_booking(self.object)

class BookingDeleteView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, DeleteView):

    model = Booking

    def delete(self, request, *args, **kwargs):
        result = super(BookingDeleteView, self).delete(request, *args, **kwargs)
        end_date = ''
        if self.object.end_date:
            end_date = ' to {}, '.format(
                self.object.end_date.strftime('%d/%m/%Y')
            )
        messages.info(
            self.request,
            "Deleted booking from {}{} {}".format(
                self.object.start_date.strftime('%d/%m/%Y'),
                end_date,
                self.object.title,
            )
        )
        return result

    def get_success_url(self):
        return reverse('booking.list')

class BookingDetailView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, DetailView):

    model = Booking

    def get_context_data(self, **kwargs):
        context = super(BookingDetailView, self).get_context_data(**kwargs)
        context.update(dict(
            booking_settings=BookingSettings.load(),
        ))
        return context

class BookingListMixin(LoginRequiredMixin, BaseMixin, ListView):

    def get_context_data(self, **kwargs):
        context = super(BookingListMixin, self).get_context_data(**kwargs)
        context.update(dict(
            booking_settings=BookingSettings.load(),
        ))
        return context

class BookingListMonthMixin(BookingListMixin):

    def _get_date(self):
        month = int(self.kwargs.get('month', 0))
        year = int(self.kwargs.get('year', 0))
        try:
            return datetime(year, month, 1).date()
        except ValueError:
            raise Http404("Invalid date.")

    def get_context_data(self, **kwargs):
        context = super(BookingListMonthMixin, self).get_context_data(**kwargs)
        d = self._get_date()
        first_next_month = d + relativedelta(months=+1, day=1)
        first_prev_month = d + relativedelta(months=-1, day=1)
        context.update(dict(
            first_next_month=first_next_month,
            first_prev_month=first_prev_month,
            sub_heading="Bookings for {}".format(d.strftime("%B %Y"))
        ))
        return context

    def get_queryset(self):
        d = self._get_date()
        return Booking.objects.month(self.request.user, d.month, d.year)

class BookingListTodayMixin(BookingListMixin):

    def get_context_data(self, **kwargs):
        context = super(BookingListTodayMixin, self).get_context_data(**kwargs)
        today = datetime.today().date()
        first_next_month = today + relativedelta(months=+1, day=1)
        first_prev_month = today + relativedelta(months=-1, day=1)
        context.update(dict(
            booking_settings=BookingSettings.load(),
            first_next_month=first_next_month,
            first_prev_month=first_prev_month,
            display_today=True,
            sub_heading="Home",
        ))
        return context

    def get_queryset(self):
        return Booking.objects.calendar(self.request.user)

class BookingListView(BookingListTodayMixin):

    def get_context_data(self, **kwargs):
        context = super(BookingListView, self).get_context_data(**kwargs)
        context.update(dict(
            booking_edit=True,
        ))
        return context

    model = Booking

class BookingListMonthView(BookingListMonthMixin):

    model = Booking

    def get_context_data(self, **kwargs):
        context = super(BookingListMonthView, self).get_context_data(**kwargs)
        context.update(dict(
            booking_edit=True,
        ))
        return context

class BookingUpdateNotesStaffView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = BookingNotesStaffForm
    model = Booking
    template_name = 'booking/booking_notes_form.html'

    def get_success_url(self):
        return _url_booking(self.object)

class BookingUpdateNotesUserView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = BookingNotesUserForm
    model = Booking
    template_name = 'booking/booking_notes_form.html'

    def get_success_url(self):
        return _url_booking(self.object)

class BookingUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    model = Booking

    def get_success_url(self):
        return _url_booking(self.object)

class CategoryCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = CategoryForm
    model = Category

    def get_success_url(self):
        return reverse('booking.category.list')

class CategoryListView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, ListView):

    model = Category

class CategoryUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = CategoryForm
    model = Category

    def get_success_url(self):
        return reverse('booking.category.list')

def _get_response(file_name):
    # Create the HttpResponse object with the appropriate CSV header.
    response = HttpResponse(content_type='application/pdf')
    response['Content-Disposition'] = 'attachment; filename="{}"'.format(file_name)
    return response

@login_required
def download_calendar(request):
    response = _get_response('calendar.pdf')
    PdfCalendar().report(response, request.user)
    return response

class LocationCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = LocationForm
    model = Location

    def get_success_url(self):
        return reverse('booking.location.list')

class LocationDetailView(BaseMixin, DetailView):

    model = Location

class LocationListView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, ListView):

    model = Location

class LocationUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = LocationForm
    model = Location

    def get_success_url(self):
        return reverse('booking.location.list')

class RotaCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = RotaForm
    model = Rota

    def _get_booking(self):
        pk = self.kwargs.get('pk', None)
        return Booking.objects.get(pk=pk)

    def get_context_data(self, **kwargs):
        context = super(RotaCreateView, self).get_context_data(**kwargs)
        context.update(dict(booking=self._get_booking()))
        return context

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.booking = self._get_booking()
        return super(RotaCreateView, self).form_valid(form)

    def get_success_url(self):
        return _url_booking(self.object.booking)

class RotaDeleteView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = RotaEmptyForm
    model = Rota
    template_name = 'booking/rota_remove_form.html'

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.deleted = True
        return super(RotaDeleteView, self).form_valid(form)

    def get_success_url(self):
        return _url_booking(self.object.booking)

class RotaUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = RotaForm
    model = Rota

    def get_context_data(self, **kwargs):
        context = super(RotaUpdateView, self).get_context_data(**kwargs)
        context.update(dict(booking=self.object.booking))
        return context

    def get_success_url(self):
        return _url_booking(self.object.booking)

class RotaTypeCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = RotaTypeForm
    model = RotaType

    def get_success_url(self):
        return reverse('booking.rota.type.list')

class RotaTypeListView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, ListView):

    model = RotaType

class RotaTypeUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = RotaTypeForm
    model = RotaType

    def get_success_url(self):
        return reverse('booking.rota.type.list')
<commit_msg>'booking' app - calendar mixin #1181<commit_after># -*- encoding: utf-8 -*-
from datetime import datetime

from dateutil.relativedelta import relativedelta

from django.contrib import messages
from django.contrib.auth.decorators import login_required
from django.core.urlresolvers import reverse
from django.http import (
    Http404,
    HttpResponse,
)
from django.views.generic import (
    CreateView,
    DeleteView,
    DetailView,
    ListView,
    UpdateView,
)

from braces.views import (
    LoginRequiredMixin,
    StaffuserRequiredMixin,
)

from base.view_utils import BaseMixin

from .forms import (
    BookingForm,
    BookingNotesStaffForm,
    BookingNotesUserForm,
    CategoryForm,
    LocationForm,
    RotaEmptyForm,
    RotaForm,
    RotaTypeForm,
)
from .models import (
    Booking,
    BookingSettings,
    Category,
    Location,
    Rota,
    RotaType,
)
from .service import (
    grouper,
    HtmlCalendar,
    PdfCalendar,
)

def _url_booking(booking):
    """If we are editing from the detail form."""
    if BookingSettings.load().edit_from_detail:
        return reverse('booking.detail', args=[booking.pk])
    else:
        return _url_booking_list_month(booking)

def _url_booking_list_month(booking):
    return reverse(
        'booking.list.month',
        kwargs=dict(
            year=booking.start_date.year,
            month=booking.start_date.month
        )
    )

class BookingCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    model = Booking

    def get_success_url(self):
        return _url_booking(self.object)

class BookingDeleteView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, DeleteView):

    model = Booking

    def delete(self, request, *args, **kwargs):
        result = super(BookingDeleteView, self).delete(request, *args, **kwargs)
        end_date = ''
        if self.object.end_date:
            end_date = ' to {}, '.format(
                self.object.end_date.strftime('%d/%m/%Y')
            )
        messages.info(
            self.request,
            "Deleted booking from {}{} {}".format(
                self.object.start_date.strftime('%d/%m/%Y'),
                end_date,
                self.object.title,
            )
        )
        return result

    def get_success_url(self):
        return reverse('booking.list')

class BookingDetailView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, DetailView):

    model = Booking

    def get_context_data(self, **kwargs):
        context = super(BookingDetailView, self).get_context_data(**kwargs)
        context.update(dict(
            booking_settings=BookingSettings.load(),
        ))
        return context

class BookingListMixin(LoginRequiredMixin, BaseMixin, ListView):

    def get_context_data(self, **kwargs):
        context = super(BookingListMixin, self).get_context_data(**kwargs)
        context.update(dict(
            booking_settings=BookingSettings.load(),
        ))
        return context

class BookingListMonthMixin(BookingListMixin):

    def _get_date(self):
        month = int(self.kwargs.get('month', 0))
        year = int(self.kwargs.get('year', 0))
        try:
            return datetime(year, month, 1).date()
        except ValueError:
            raise Http404("Invalid date.")

    def get_context_data(self, **kwargs):
        context = super(BookingListMonthMixin, self).get_context_data(**kwargs)
        d = self._get_date()
        first_next_month = d + relativedelta(months=+1, day=1)
        first_prev_month = d + relativedelta(months=-1, day=1)
        context.update(dict(
            first_next_month=first_next_month,
            first_prev_month=first_prev_month,
            sub_heading="Bookings for {}".format(d.strftime("%B %Y"))
        ))
        return context

    def get_queryset(self):
        d = self._get_date()
        return Booking.objects.month(self.request.user, d.month, d.year)

class BookingListTodayMixin(BookingListMixin):

    def get_context_data(self, **kwargs):
        context = super(BookingListTodayMixin, self).get_context_data(**kwargs)
        today = datetime.today().date()
        first_next_month = today + relativedelta(months=+1, day=1)
        first_prev_month = today + relativedelta(months=-1, day=1)
        context.update(dict(
            booking_settings=BookingSettings.load(),
            first_next_month=first_next_month,
            first_prev_month=first_prev_month,
            display_today=True,
            sub_heading="Home",
        ))
        return context

    def get_queryset(self):
        return Booking.objects.calendar(self.request.user)

class BookingListView(BookingListTodayMixin):

    def get_context_data(self, **kwargs):
        context = super(BookingListView, self).get_context_data(**kwargs)
        context.update(dict(
            booking_edit=True,
        ))
        return context

    model = Booking

class BookingListMonthView(BookingListMonthMixin):

    model = Booking

    def get_context_data(self, **kwargs):
        context = super(BookingListMonthView, self).get_context_data(**kwargs)
        context.update(dict(
            booking_edit=True,
        ))
        return context

class BookingUpdateNotesStaffView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = BookingNotesStaffForm
    model = Booking
    template_name = 'booking/booking_notes_form.html'

    def get_success_url(self):
        return _url_booking(self.object)

class BookingUpdateNotesUserView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = BookingNotesUserForm
    model = Booking
    template_name = 'booking/booking_notes_form.html'

    def get_success_url(self):
        return _url_booking(self.object)

class BookingUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    model = Booking

    def get_success_url(self):
        return _url_booking(self.object)

class CalendarMixin(object):

    def get_context_data(self, **kwargs):
        context = super().get_context_data(**kwargs)
        calendars = HtmlCalendar().get_calendars()
        grouped = grouper(calendars, 3)
        context.update(dict(
            calendar=list(grouped),
        ))
        return context

class CategoryCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = CategoryForm
    model = Category

    def get_success_url(self):
        return reverse('booking.category.list')

class CategoryListView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, ListView):

    model = Category

class CategoryUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = CategoryForm
    model = Category

    def get_success_url(self):
        return reverse('booking.category.list')

def _get_response(file_name):
    # Create the HttpResponse object with the appropriate CSV header.
    response = HttpResponse(content_type='application/pdf')
    response['Content-Disposition'] = 'attachment; filename="{}"'.format(file_name)
    return response

@login_required
def download_calendar(request):
    response = _get_response('calendar.pdf')
    PdfCalendar().report(response, request.user)
    return response

class LocationCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = LocationForm
    model = Location

    def get_success_url(self):
        return reverse('booking.location.list')

class LocationDetailView(BaseMixin, DetailView):

    model = Location

class LocationListView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, ListView):

    model = Location

class LocationUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = LocationForm
    model = Location

    def get_success_url(self):
        return reverse('booking.location.list')

class RotaCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = RotaForm
    model = Rota

    def _get_booking(self):
        pk = self.kwargs.get('pk', None)
        return Booking.objects.get(pk=pk)

    def get_context_data(self, **kwargs):
        context = super(RotaCreateView, self).get_context_data(**kwargs)
        context.update(dict(booking=self._get_booking()))
        return context

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.booking = self._get_booking()
        return super(RotaCreateView, self).form_valid(form)

    def get_success_url(self):
        return _url_booking(self.object.booking)

class RotaDeleteView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = RotaEmptyForm
    model = Rota
    template_name = 'booking/rota_remove_form.html'

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.deleted = True
        return super(RotaDeleteView, self).form_valid(form)

    def get_success_url(self):
        return _url_booking(self.object.booking)

class RotaUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = RotaForm
    model = Rota

    def get_context_data(self, **kwargs):
        context = super(RotaUpdateView, self).get_context_data(**kwargs)
        context.update(dict(booking=self.object.booking))
        return context

    def get_success_url(self):
        return _url_booking(self.object.booking)

class RotaTypeCreateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, CreateView):

    form_class = RotaTypeForm
    model = RotaType

    def get_success_url(self):
        return reverse('booking.rota.type.list')

class RotaTypeListView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, ListView):

    model = RotaType

class RotaTypeUpdateView(
        LoginRequiredMixin, StaffuserRequiredMixin, BaseMixin, UpdateView):

    form_class = RotaTypeForm
    model = RotaType

    def get_success_url(self):
        return reverse('booking.rota.type.list')
<|endoftext|>


<commit_before>from corehq.apps.app_manager import id_strings
from corehq.apps.app_manager.suite_xml.contributors import (
    SuiteContributorByModule,
)
from corehq.apps.app_manager.suite_xml.post_process.instances import (
    get_all_instances_referenced_in_xpaths,
)
from corehq.apps.app_manager.suite_xml.sections.details import (
    DetailsHelper,
    get_instances_for_module,
)
from corehq.apps.app_manager.suite_xml.xml_models import (
    Command,
    Display,
    Instance,
    Itemset,
    PushFrame,
    QueryData,
    QueryPrompt,
    RemoteRequest,
    RemoteRequestPost,
    RemoteRequestQuery,
    RemoteRequestSession,
    SessionDatum,
    Stack,
    Text,
)
from corehq.apps.app_manager.util import module_offers_search
from corehq.apps.app_manager.xpath import CaseTypeXpath, InstanceXpath, interpolate_xpath
from corehq.apps.case_search.models import CASE_SEARCH_BLACKLISTED_OWNER_ID_KEY
from corehq.util.view_utils import absolute_reverse

RESULTS_INSTANCE = 'results'  # The name of the instance where search results are stored
SESSION_INSTANCE = 'commcaresession'

class QuerySessionXPath(InstanceXpath):
    id = SESSION_INSTANCE

    @property
    def path(self):
        return 'session/data/{}'.format(self)

class RemoteRequestFactory(object):
    def __init__(self, domain, app, module):
        self.domain = domain
        self.app = app
        self.module = module

    def build_remote_request(self):
        return RemoteRequest(
            post=self._build_remote_request_post(),
            command=self._build_command(),
            instances=self._build_instances(),
            session=self._build_session(),
            stack=self._build_stack(),
        )

    def _build_remote_request_post(self):
        return RemoteRequestPost(
            url=absolute_reverse('claim_case', args=[self.domain]),
            relevant=self.module.search_config.relevant,
            data=[
                QueryData(
                    key='case_id',
                    ref=QuerySessionXPath('case_id').instance(),
                ),
            ]
        )

    def _build_command(self):
        return Command(
            id=id_strings.search_command(self.module),
            display=Display(
                text=Text(locale_id=id_strings.case_search_locale(self.module)),
            ),
        )

    def _build_instances(self):
        prompt_select_instances = [
            Instance(id=prop.itemset.instance_id, src=prop.itemset.instance_uri)
            for prop in self.module.search_config.properties
            if prop.itemset.instance_id
        ]

        query_xpaths = [datum.ref for datum in self._get_remote_request_query_datums()]
        config_xpaths = [self.module.search_config.relevant, self.module.search_config.search_filter]
        instances, unknown_instances = get_all_instances_referenced_in_xpaths(
            self.app,
            query_xpaths + config_xpaths
        )
        # we use the module's case list/details view to select the datum so also
        # need these instances to be available
        instances |= get_instances_for_module(self.app, self.module)

        # sorted list to prevent intermittent test failures
        return sorted(list(instances) + prompt_select_instances, key=lambda i: i.id)

    def _build_session(self):
        return RemoteRequestSession(
            queries=self._build_remote_request_queries(),
            data=self._build_remote_request_datums(),
        )

    def _build_remote_request_queries(self):
        return [
            RemoteRequestQuery(
                url=absolute_reverse('remote_search', args=[self.app.domain]),
                storage_instance=RESULTS_INSTANCE,
                template='case',
                data=self._get_remote_request_query_datums(),
                prompts=self._build_query_prompts()
            )
        ]

    def _build_remote_request_datums(self):
        details_helper = DetailsHelper(self.app)
        if self.module.case_details.short.custom_xml:
            short_detail_id = 'case_short'
        else:
            short_detail_id = 'search_short'

        nodeset = CaseTypeXpath(self.module.case_type).case(instance_name=RESULTS_INSTANCE)
        if self.module.search_config.search_filter:
            nodeset = f"{nodeset}[{interpolate_xpath(self.module.search_config.search_filter)}]"

        return [SessionDatum(
            id='case_id',
            nodeset=nodeset,
            value='./@case_id',
            detail_select=details_helper.get_detail_id_safe(self.module, short_detail_id),
            detail_confirm=details_helper.get_detail_id_safe(self.module, 'case_long'),
        )]

    def _get_remote_request_query_datums(self):
        default_query_datums = [
            QueryData(
                key='case_type',
                ref="'{}'".format(self.module.case_type)
            ),
            QueryData(
                key='include_closed',
                ref="'{}'".format(self.module.search_config.include_closed)
            )
        ]
        extra_query_datums = [
            QueryData(key="{}".format(c.property), ref="{}".format(c.defaultValue))
            for c in self.module.search_config.default_properties
        ]
        if self.module.search_config.blacklisted_owner_ids_expression:
            extra_query_datums.append(
                QueryData(
                    key=CASE_SEARCH_BLACKLISTED_OWNER_ID_KEY,
                    ref="{}".format(self.module.search_config.blacklisted_owner_ids_expression)
                )
            )
        return default_query_datums + extra_query_datums

    def _build_query_prompts(self):
        prompts = []
        for prop in self.module.search_config.properties:
            kwargs = {
                'key': prop.name,
                'display': Display(text=Text(locale_id=id_strings.search_property_locale(self.module, prop.name))),
            }
            if prop.appearance and self.app.enable_search_prompt_appearance:
                kwargs['appearance'] = prop.appearance
            if prop.input_:
                kwargs['input_'] = prop.input_
            if prop.default_value and self.app.enable_default_value_expression:
                kwargs['default_value'] = prop.default_value
            if prop.itemset.nodeset:
                kwargs['itemset'] = Itemset(
                    nodeset=prop.itemset.nodeset,
                    label_ref=prop.itemset.label,
                    value_ref=prop.itemset.value,
                    sort_ref=prop.itemset.sort,
                )
            prompts.append(QueryPrompt(**kwargs))
        return prompts

    def _build_stack(self):
        stack = Stack()
        frame = PushFrame()
        frame.add_rewind(QuerySessionXPath('case_id').instance())
        stack.add_frame(frame)
        return stack

class RemoteRequestContributor(SuiteContributorByModule):
    """
    Adds a remote-request node, which sets the URL and query details for
    synchronous searching and case claiming.

    Search is available from the module's case list.

    See "remote-request" in the `CommCare 2.0 Suite Definition`_ for details.

    .. _CommCare 2.0 Suite Definition: https://github.com/dimagi/commcare/wiki/Suite20#remote-request

    """
    def get_module_contributions(self, module):
        if module_offers_search(module):
            return [RemoteRequestFactory(self.app.domain, self.app, module).build_remote_request()]
        return []
<commit_msg>add instances referenced in the default value expressions<commit_after>from corehq.apps.app_manager import id_strings
from corehq.apps.app_manager.suite_xml.contributors import (
    SuiteContributorByModule,
)
from corehq.apps.app_manager.suite_xml.post_process.instances import (
    get_all_instances_referenced_in_xpaths,
)
from corehq.apps.app_manager.suite_xml.sections.details import (
    DetailsHelper,
    get_instances_for_module,
)
from corehq.apps.app_manager.suite_xml.xml_models import (
    Command,
    Display,
    Instance,
    Itemset,
    PushFrame,
    QueryData,
    QueryPrompt,
    RemoteRequest,
    RemoteRequestPost,
    RemoteRequestQuery,
    RemoteRequestSession,
    SessionDatum,
    Stack,
    Text,
)
from corehq.apps.app_manager.util import module_offers_search
from corehq.apps.app_manager.xpath import CaseTypeXpath, InstanceXpath, interpolate_xpath
from corehq.apps.case_search.models import CASE_SEARCH_BLACKLISTED_OWNER_ID_KEY
from corehq.util.view_utils import absolute_reverse

RESULTS_INSTANCE = 'results'  # The name of the instance where search results are stored
SESSION_INSTANCE = 'commcaresession'

class QuerySessionXPath(InstanceXpath):
    id = SESSION_INSTANCE

    @property
    def path(self):
        return 'session/data/{}'.format(self)

class RemoteRequestFactory(object):
    def __init__(self, domain, app, module):
        self.domain = domain
        self.app = app
        self.module = module

    def build_remote_request(self):
        return RemoteRequest(
            post=self._build_remote_request_post(),
            command=self._build_command(),
            instances=self._build_instances(),
            session=self._build_session(),
            stack=self._build_stack(),
        )

    def _build_remote_request_post(self):
        return RemoteRequestPost(
            url=absolute_reverse('claim_case', args=[self.domain]),
            relevant=self.module.search_config.relevant,
            data=[
                QueryData(
                    key='case_id',
                    ref=QuerySessionXPath('case_id').instance(),
                ),
            ]
        )

    def _build_command(self):
        return Command(
            id=id_strings.search_command(self.module),
            display=Display(
                text=Text(locale_id=id_strings.case_search_locale(self.module)),
            ),
        )

    def _build_instances(self):
        prompt_select_instances = [
            Instance(id=prop.itemset.instance_id, src=prop.itemset.instance_uri)
            for prop in self.module.search_config.properties
            if prop.itemset.instance_id
        ]

        query_xpaths = [datum.ref for datum in self._get_remote_request_query_datums()]
        config_xpaths = [self.module.search_config.relevant, self.module.search_config.search_filter]
        query_xpaths.extend([prop.default_value for prop in self.module.search_config.properties])
        instances, unknown_instances = get_all_instances_referenced_in_xpaths(
            self.app,
            query_xpaths + config_xpaths
        )
        # we use the module's case list/details view to select the datum so also
        # need these instances to be available
        instances |= get_instances_for_module(self.app, self.module)

        # sorted list to prevent intermittent test failures
        return sorted(list(instances) + prompt_select_instances, key=lambda i: i.id)

    def _build_session(self):
        return RemoteRequestSession(
            queries=self._build_remote_request_queries(),
            data=self._build_remote_request_datums(),
        )

    def _build_remote_request_queries(self):
        return [
            RemoteRequestQuery(
                url=absolute_reverse('remote_search', args=[self.app.domain]),
                storage_instance=RESULTS_INSTANCE,
                template='case',
                data=self._get_remote_request_query_datums(),
                prompts=self._build_query_prompts()
            )
        ]

    def _build_remote_request_datums(self):
        details_helper = DetailsHelper(self.app)
        if self.module.case_details.short.custom_xml:
            short_detail_id = 'case_short'
        else:
            short_detail_id = 'search_short'

        nodeset = CaseTypeXpath(self.module.case_type).case(instance_name=RESULTS_INSTANCE)
        if self.module.search_config.search_filter:
            nodeset = f"{nodeset}[{interpolate_xpath(self.module.search_config.search_filter)}]"

        return [SessionDatum(
            id='case_id',
            nodeset=nodeset,
            value='./@case_id',
            detail_select=details_helper.get_detail_id_safe(self.module, short_detail_id),
            detail_confirm=details_helper.get_detail_id_safe(self.module, 'case_long'),
        )]

    def _get_remote_request_query_datums(self):
        default_query_datums = [
            QueryData(
                key='case_type',
                ref="'{}'".format(self.module.case_type)
            ),
            QueryData(
                key='include_closed',
                ref="'{}'".format(self.module.search_config.include_closed)
            )
        ]
        extra_query_datums = [
            QueryData(key="{}".format(c.property), ref="{}".format(c.defaultValue))
            for c in self.module.search_config.default_properties
        ]
        if self.module.search_config.blacklisted_owner_ids_expression:
            extra_query_datums.append(
                QueryData(
                    key=CASE_SEARCH_BLACKLISTED_OWNER_ID_KEY,
                    ref="{}".format(self.module.search_config.blacklisted_owner_ids_expression)
                )
            )
        return default_query_datums + extra_query_datums

    def _build_query_prompts(self):
        prompts = []
        for prop in self.module.search_config.properties:
            kwargs = {
                'key': prop.name,
                'display': Display(text=Text(locale_id=id_strings.search_property_locale(self.module, prop.name))),
            }
            if prop.appearance and self.app.enable_search_prompt_appearance:
                kwargs['appearance'] = prop.appearance
            if prop.input_:
                kwargs['input_'] = prop.input_
            if prop.default_value and self.app.enable_default_value_expression:
                kwargs['default_value'] = prop.default_value
            if prop.itemset.nodeset:
                kwargs['itemset'] = Itemset(
                    nodeset=prop.itemset.nodeset,
                    label_ref=prop.itemset.label,
                    value_ref=prop.itemset.value,
                    sort_ref=prop.itemset.sort,
                )
            prompts.append(QueryPrompt(**kwargs))
        return prompts

    def _build_stack(self):
        stack = Stack()
        frame = PushFrame()
        frame.add_rewind(QuerySessionXPath('case_id').instance())
        stack.add_frame(frame)
        return stack

class RemoteRequestContributor(SuiteContributorByModule):
    """
    Adds a remote-request node, which sets the URL and query details for
    synchronous searching and case claiming.

    Search is available from the module's case list.

    See "remote-request" in the `CommCare 2.0 Suite Definition`_ for details.

    .. _CommCare 2.0 Suite Definition: https://github.com/dimagi/commcare/wiki/Suite20#remote-request

    """
    def get_module_contributions(self, module):
        if module_offers_search(module):
            return [RemoteRequestFactory(self.app.domain, self.app, module).build_remote_request()]
        return []
<|endoftext|>


<commit_before>from rctk.widgets.control import Control, remote_attribute

from rctk.task import Task
from rctk.event import Changable, Submittable

class Text(Control, Changable, Submittable):
    name = "text"

    value = remote_attribute('value', "")

    def __init__(self, tk, value="", rows=1, columns=20):
        self._value = value
        self._rows = rows
        self._columns = columns
        super(Text, self).__init__(tk)

    def create(self):
        self.tk.create_control(self, value=self._value, rows=self._rows, columns=self._columns)

    def sync(self, **data):
        if 'value' in data:
            self._value = data['value']

class Password(Text):
    name = "password"
<commit_msg>Allow additional properties on Text<commit_after>from rctk.widgets.control import Control, remote_attribute

from rctk.task import Task
from rctk.event import Changable, Submittable

class Text(Control, Changable, Submittable):
    name = "text"

    value = remote_attribute('value', "")

    def __init__(self, tk, value="", rows=1, columns=20, **properties):
        self._value = value
        self._rows = rows
        self._columns = columns
        super(Text, self).__init__(tk, **properties)

    def create(self):
        self.tk.create_control(self, value=self._value, rows=self._rows, columns=self._columns)

    def sync(self, **data):
        if 'value' in data:
            self._value = data['value']

class Password(Text):
    name = "password"
<|endoftext|>


<commit_before>"""
Class for the SDO/AIA instrument. Holds information about the cadence and
spatial and spectroscopic resolution.
"""

import os
import sys
import logging
from collections import namedtuple

import numpy as np
import scipy.interpolate
import astropy.units as u
import sunpy.map

Pair = namedtuple('Pair','x y')

class InstrumentSDOAIA(object):
    """
    Atmospheric Imaging Assembly object for observing synthesized active
    region emission.

    Parameters
    ----------
    observing_time : `tuple`
        start and end of observing time
    observing_area : `tuple`
        x and y range of observation
    use_temperature_response_functions : `bool`
        if True, do simple counts calculation
    response_function_file : `str`
        filename containing AIA response functions

    Examples
    --------
    Notes
    -----
    """

    fits_template = sunpy.map.header.MapMeta()
    fits_template['telescop'] = 'SDO/AIA'
    fits_template['detector'] = 'AIA'
    fits_template['waveunit'] = 'angstrom'

    name = 'SDO_AIA'
    channels = [{'wavelength':94*u.angstrom,'telescope_number':4},
                {'wavelength':131*u.angstrom,'telescope_number':1},
                {'wavelength':171*u.angstrom,'telescope_number':3},
                {'wavelength':193*u.angstrom,'telescope_number':2},
                {'wavelength':211*u.angstrom,'telescope_number':2},
                {'wavelength':335*u.angstrom,'telescope_number':1}]

    cadence = 10.0*u.s
    resolution = Pair(0.600698*u.arcsec/u.pixel,0.600698*u.arcsec/u.pixel)

    def __init__(self, observing_time, observing_area=None,
    use_temperature_response_functions=True,response_function_file=''):
        """
        Constructor
        """
        self.logger = logging.getLogger(name=type(self).__name__)
        self.observing_time = np.arange(
                            observing_time[0].to(u.s).value,
                            observing_time[1].to(u.s).value,
                            self.cadence.value)*u.s
        self.observing_area = observing_area
        self.use_temperature_response_functions = use_temperature_response_functions
        if self.use_temperature_response_functions and response_function_file:
            self._setup_response_functions(response_function_file)

    def detect(self,loop,channel):
        """
        For a given loop object, calculate the counts detected by AIA in a
        particular channel.

        Parameters
        ----------
        loop : loop object
        channel : `dict`

        Returns
        -------
        counts : array-like
        """
        if self.use_temperature_response_functions:
            counts = self._detect_simple(loop,channel)
        else:
            counts = self._detect_full(loop,channel)
        return counts

    def _detect_simple(self,loop,channel):
        """
        Calculate counts using the density and temperature response functions.
        No emissivity model needed.
        """
        response_function = channel['response_interpolator'](np.ravel(loop.temperature))*u.counts*u.cm**5/u.s/u.pixel

        return np.reshape(np.ravel(loop.temperature**2)*response_function,
                          np.shape(loop.density))

    def _detect_full(self,loop,channel):
        """
        Calculate counts use emissivity for a large number of transitions.
        Requires emissivity model.

        Notes
        -----
        This is necessary when taking into account non-equilibrium ionization.
        """
        pass

    def _setup_response_functions(self,filename):
        """
        Setup interpolators from the AIA temperature response functions.

        Notes
        -----
        This should be replaced once the response functions are available in
        SunPy.
        Probably should configure wavelength response function
        interpolators also.
        """
        _tmp = np.loadtxt(filename)
        channel_order = {c:i for c,i in zip([94,131,171,
                                             193,211,335]*u.angstrom,range(6))}
        _tmp_temperature = 10**(_tmp[:,0])
        for i,channel in enumerate(self.channels):
            _tmp_response = _tmp[:,channel_order[channel['wavelength']]+1]
            self.channels[i]['response_interpolator'] = interpolate.interp1d(_tmp_temperature,_tmp_response)

    def make_fits_header(self,field,channel):
        """
        Build up FITS header with relevant instrument information.
        """
        update_entries = ['crpix1','crpix2','crval1','crval2','cunit1',
                          'cunit2','crlt_obs','ctype1','ctype2','date-obs',
                          'dsun_obs','rsun_obs']
        fits_header = self.fits_template.copy()
        for entry in update_entries:
            fits_header[entry] = field.clipped_hmi_map.meta[entry]
        fits_header['cdelt1'] = self.resolution.x.value
        fits_header['cdelt2'] = self.resolution.y.value
        fits_header['instrume'] = 'AIA_' + channel['telescope_number']
        fits_header['wavelnth'] = int(channel['wavelength'])

        return fits_header

    def make_detector_array(self,field):
        """
        Construct bins based on desired observing area.
        """
        delta_x = np.fabs(field.clipped_hmi_map.xrange[1] -
                          field.clipped_hmi_map.xrange[0])
        delta_y = np.fabs(field.clipped_hmi_map.yrange[1] -
                          field.clipped_hmi_map.yrange[0])
        self.bins = [int(np.ceil(delta_x/self.resolution.x)).value,
                     int(np.ceil(delta_y/self.resolution.y)).value]
        self.bin_ranges = [
            field._convert_angle_to_length(field.clipped_hmi_map.xrange).value,
            field._convert_angle_to_length(field.clipped_hmi_map.yrange).value]
<commit_msg>more interpolate fixes<commit_after>"""
Class for the SDO/AIA instrument. Holds information about the cadence and
spatial and spectroscopic resolution.
"""

import os
import sys
import logging
from collections import namedtuple

import numpy as np
import scipy.interpolate
import astropy.units as u
import sunpy.map

Pair = namedtuple('Pair','x y')

class InstrumentSDOAIA(object):
    """
    Atmospheric Imaging Assembly object for observing synthesized active
    region emission.

    Parameters
    ----------
    observing_time : `tuple`
        start and end of observing time
    observing_area : `tuple`
        x and y range of observation
    use_temperature_response_functions : `bool`
        if True, do simple counts calculation
    response_function_file : `str`
        filename containing AIA response functions

    Examples
    --------
    Notes
    -----
    """

    fits_template = sunpy.map.header.MapMeta()
    fits_template['telescop'] = 'SDO/AIA'
    fits_template['detector'] = 'AIA'
    fits_template['waveunit'] = 'angstrom'

    name = 'SDO_AIA'
    channels = [{'wavelength':94*u.angstrom,'telescope_number':4},
                {'wavelength':131*u.angstrom,'telescope_number':1},
                {'wavelength':171*u.angstrom,'telescope_number':3},
                {'wavelength':193*u.angstrom,'telescope_number':2},
                {'wavelength':211*u.angstrom,'telescope_number':2},
                {'wavelength':335*u.angstrom,'telescope_number':1}]

    cadence = 10.0*u.s
    resolution = Pair(0.600698*u.arcsec/u.pixel,0.600698*u.arcsec/u.pixel)

    def __init__(self, observing_time, observing_area=None,
    use_temperature_response_functions=True,response_function_file=''):
        """
        Constructor
        """
        self.logger = logging.getLogger(name=type(self).__name__)
        self.observing_time = np.arange(
                            observing_time[0].to(u.s).value,
                            observing_time[1].to(u.s).value,
                            self.cadence.value)*u.s
        self.observing_area = observing_area
        self.use_temperature_response_functions = use_temperature_response_functions
        if self.use_temperature_response_functions and response_function_file:
            self._setup_response_functions(response_function_file)

    def detect(self,loop,channel):
        """
        For a given loop object, calculate the counts detected by AIA in a
        particular channel.

        Parameters
        ----------
        loop : loop object
        channel : `dict`

        Returns
        -------
        counts : array-like
        """
        if self.use_temperature_response_functions:
            counts = self._detect_simple(loop,channel)
        else:
            counts = self._detect_full(loop,channel)
        return counts

    def _detect_simple(self,loop,channel):
        """
        Calculate counts using the density and temperature response functions.
        No emissivity model needed.
        """
        response_function = channel['response_interpolator'](np.ravel(loop.temperature))*u.counts*u.cm**5/u.s/u.pixel

        return np.reshape(np.ravel(loop.temperature**2)*response_function,
                          np.shape(loop.density))

    def _detect_full(self,loop,channel):
        """
        Calculate counts use emissivity for a large number of transitions.
        Requires emissivity model.

        Notes
        -----
        This is necessary when taking into account non-equilibrium ionization.
        """
        pass

    def _setup_response_functions(self,filename):
        """
        Setup interpolators from the AIA temperature response functions.

        Notes
        -----
        This should be replaced once the response functions are available in
        SunPy.
        Probably should configure wavelength response function
        interpolators also.
        """
        _tmp = np.loadtxt(filename)
        channel_order = {c:i for c,i in zip([94,131,171,
                                             193,211,335]*u.angstrom,range(6))}
        _tmp_temperature = 10**(_tmp[:,0])
        for i,channel in enumerate(self.channels):
            _tmp_response = _tmp[:,channel_order[channel['wavelength']]+1]
            self.channels[i]['response_interpolator'] = scipy.interpolate.interp1d(_tmp_temperature,_tmp_response)

    def make_fits_header(self,field,channel):
        """
        Build up FITS header with relevant instrument information.
        """
        update_entries = ['crpix1','crpix2','crval1','crval2','cunit1',
                          'cunit2','crlt_obs','ctype1','ctype2','date-obs',
                          'dsun_obs','rsun_obs']
        fits_header = self.fits_template.copy()
        for entry in update_entries:
            fits_header[entry] = field.clipped_hmi_map.meta[entry]
        fits_header['cdelt1'] = self.resolution.x.value
        fits_header['cdelt2'] = self.resolution.y.value
        fits_header['instrume'] = 'AIA_' + channel['telescope_number']
        fits_header['wavelnth'] = int(channel['wavelength'])

        return fits_header

    def make_detector_array(self,field):
        """
        Construct bins based on desired observing area.
        """
        delta_x = np.fabs(field.clipped_hmi_map.xrange[1] -
                          field.clipped_hmi_map.xrange[0])
        delta_y = np.fabs(field.clipped_hmi_map.yrange[1] -
                          field.clipped_hmi_map.yrange[0])
        self.bins = [int(np.ceil(delta_x/self.resolution.x)).value,
                     int(np.ceil(delta_y/self.resolution.y)).value]
        self.bin_ranges = [
            field._convert_angle_to_length(field.clipped_hmi_map.xrange).value,
            field._convert_angle_to_length(field.clipped_hmi_map.yrange).value]
<|endoftext|>