<commit_before>#! usr/bin/python3
# -*- coding: utf8 -*-
#
# Flicket - copyright Paul Bourne: evereux@gmail.com

import datetime

from flask import render_template, redirect, url_for, g, request, flash
from flask_login import login_required

from . import flicket_bp
from application import app, db
from application.flicket.forms.flicket_forms import ReplyForm
from application.flicket.models.flicket_models import FlicketTicket, FlicketStatus, FlicketPost
from application.flicket.scripts.flicket_functions import block_quoter
from application.flicket.scripts.flicket_upload import upload_documents, add_upload_to_db
from application.flicket.scripts.email import FlicketMail, get_recipients

# view ticket details
@flicket_bp.route(app.config['FLICKET'] + 'ticket_view/<ticket_id>/', methods=['GET', 'POST'])
@flicket_bp.route(app.config['FLICKET'] + 'ticket_view/<ticket_id>/<int:page>/', methods=['GET', 'POST'])
@login_required
def ticket_view(ticket_id, page=1):
    # todo: make sure underscores aren't allowed in usernames as it breaks markdown.

    # is ticket number legitimate
    ticket = FlicketTicket.query.filter_by(id=ticket_id).first()

    if not ticket:
        flash('Cannot find ticket: "{}"'.format(ticket_id), category='warning')
        return redirect(url_for('flicket_bp.tickets_main'))

    # find all replies to ticket.
    replies = FlicketPost.query.filter_by(ticket_id=ticket_id).order_by(FlicketPost.date_added.asc())

    post_id = request.args.get('post_id')
    ticket_rid = request.args.get('ticket_rid')

    form = ReplyForm()

    # add reply post
    if form.validate_on_submit():

        # upload file if user has selected one and the file is in accepted list of
        files = request.files.getlist("file")

        new_files = upload_documents(files)

        if new_files is False:
            flash('There was a problem uploading files.', category='danger')
            return redirect(url_for('flicket_bp.tickets_main'))

        new_reply = FlicketPost(
            ticket=ticket,
            user=g.user,
            date_added=datetime.datetime.now(),
            content=form.content.data
        )

        # add documents to database
        post_type = 'Post'
        add_upload_to_db(new_files, new_reply, post_type)

        db.session.add(new_reply)

        # change ticket status to open.
        open = FlicketStatus.query.filter_by(status='Open').first()
        ticket.current_status = open
        db.session.commit()

        # send email notification
        mail = FlicketMail()
        mail.reply_ticket(ticket=ticket)

        flash('You have replied to ticket {}: {}.'.format(ticket.id_zfill, ticket.title), category="success")
        return redirect(url_for('flicket_bp.ticket_view', ticket_id=ticket_id))

    # get post id and populate contents for auto quoting
    if post_id:
        query = FlicketPost.query.filter_by(id=post_id).first()
        reply_contents = "{} wrote on {}\r\n\r\n{}".format(query.user.name, query.date_added, query.content)
        form.content.data = block_quoter(reply_contents)
    if ticket_rid:
        reply_contents = "{} wrote on {}\r\n\r\n{}".format(ticket.user.name, ticket.date_added, ticket.content)
        form.content.data = block_quoter(reply_contents)

    replies = replies.paginate(page, app.config['posts_per_page'])

    return render_template('flicket_view.html',
                           title='Flicket - View Ticket',
                           ticket=ticket,
                           form=form,
                           replies=replies,
                           post_id=post_id,
                           page=page)
<commit_msg>Ticket status changed to 'open' if already 'closed' and a user replies to ticket.<commit_after>#! usr/bin/python3
# -*- coding: utf8 -*-
#
# Flicket - copyright Paul Bourne: evereux@gmail.com

import datetime

from flask import render_template, redirect, url_for, g, request, flash
from flask_login import login_required

from . import flicket_bp
from application import app, db
from application.flicket.forms.flicket_forms import ReplyForm
from application.flicket.models.flicket_models import FlicketTicket, FlicketStatus, FlicketPost
from application.flicket.scripts.flicket_functions import block_quoter
from application.flicket.scripts.flicket_upload import upload_documents, add_upload_to_db
from application.flicket.scripts.email import FlicketMail, get_recipients

# view ticket details
@flicket_bp.route(app.config['FLICKET'] + 'ticket_view/<ticket_id>/', methods=['GET', 'POST'])
@flicket_bp.route(app.config['FLICKET'] + 'ticket_view/<ticket_id>/<int:page>/', methods=['GET', 'POST'])
@login_required
def ticket_view(ticket_id, page=1):
    # todo: make sure underscores aren't allowed in usernames as it breaks markdown.

    # is ticket number legitimate
    ticket = FlicketTicket.query.filter_by(id=ticket_id).first()

    if not ticket:
        flash('Cannot find ticket: "{}"'.format(ticket_id), category='warning')
        return redirect(url_for('flicket_bp.tickets_main'))

    # find all replies to ticket.
    replies = FlicketPost.query.filter_by(ticket_id=ticket_id).order_by(FlicketPost.date_added.asc())

    post_id = request.args.get('post_id')
    ticket_rid = request.args.get('ticket_rid')

    form = ReplyForm()

    # add reply post
    if form.validate_on_submit():

        # upload file if user has selected one and the file is in accepted list of
        files = request.files.getlist("file")

        new_files = upload_documents(files)

        if new_files is False:
            flash('There was a problem uploading files.', category='danger')
            return redirect(url_for('flicket_bp.tickets_main'))

        new_reply = FlicketPost(
            ticket=ticket,
            user=g.user,
            date_added=datetime.datetime.now(),
            content=form.content.data
        )

        # add documents to database
        post_type = 'Post'
        add_upload_to_db(new_files, new_reply, post_type)

        db.session.add(new_reply)

        # change ticket status to open if closed.
        if ticket.current_status.status.lower() == 'closed':
            ticket_open = FlicketStatus.query.filter_by(status='Open').first()
            ticket.current_status = ticket_open

        db.session.commit()

        # send email notification
        mail = FlicketMail()
        mail.reply_ticket(ticket=ticket)

        flash('You have replied to ticket {}: {}.'.format(ticket.id_zfill, ticket.title), category="success")
        return redirect(url_for('flicket_bp.ticket_view', ticket_id=ticket_id))

    # get post id and populate contents for auto quoting
    if post_id:
        query = FlicketPost.query.filter_by(id=post_id).first()
        reply_contents = "{} wrote on {}\r\n\r\n{}".format(query.user.name, query.date_added, query.content)
        form.content.data = block_quoter(reply_contents)
    if ticket_rid:
        reply_contents = "{} wrote on {}\r\n\r\n{}".format(ticket.user.name, ticket.date_added, ticket.content)
        form.content.data = block_quoter(reply_contents)

    replies = replies.paginate(page, app.config['posts_per_page'])

    return render_template('flicket_view.html',
                           title='Flicket - View Ticket',
                           ticket=ticket,
                           form=form,
                           replies=replies,
                           post_id=post_id,
                           page=page)
<|endoftext|>


<commit_before>import os, re, urllib2

scrapeUrl = "http://interfacelift.com/wallpaper/downloads/date/widescreen_16:9/1920x1080/"
host = 'http://interfacelift.com'
saveDir = 'scraped'

# how many you want to download, 0 to download all
goal = 0
done = 0
# go to your scrape url, note down how many pages there are. this should not exceed the max.
maxPages = 209

# hi guy from the future, if this script doesn't work, you probably need to update the regex.
tag = re.compile(r"<a href=\"(?P<url>.+)\"><img.+?src=\"/img_NEW/button_download")

if not os.path.exists(saveDir) :
    os.makedirs(saveDir)

# pointless, just makes the output a little nicer.
# formats goal if 0 to read 'about XX'
goalStr = str(goal)
if goal == 0:
    # round up maxPages*10 to next hundredth
    # shitty convoluted code, i know, but works
    goal = maxPages*10 + 99
    goal = goal - (goal % 100)
    goalStr = 'about %s' % str(goal)

for page in range(1, maxPages+1):
    print 'opening page %s of %s' % (page, maxPages)

    url = scrapeUrl + "index" + str(page)

    # interfacelift returns a 403 forbidden unless you include a referer.
    headers = { 'User-agent' : "Mozilla/4.0 (compatible; MSIE 5.5; Windows NT)",
                'Referer': url}
    try:
        req = urllib2.Request(url, None, headers)
        f = urllib2.urlopen(req)

    except IOError, e:
        print "Failed to open", url
        if hasattr(e, 'code'):
            print "Error code:", e.code
        continue

    html = f.read()

    links = tag.finditer(html)

    for i in links:
        if goal == 0 or done < goal :
            print "scraping " + str(done + 1) + " of " + goalStr
            dlurl = host + i.group('url')
            req = urllib2.Request(dlurl, None, headers)
            saveFile = os.path.join(saveDir, str(done)+'.jpg')
            with open(saveFile, 'wb') as f:
                try:
                    res = urllib2.urlopen(req)
                    f.write(res.read())
                except:
                    print 'failed'
                    try:
                        os.remove(saveFile)
                    except: pass

            done += 1
<commit_msg>Fixed file deletion upon failed download<commit_after>import os, re, urllib2

scrapeUrl = 'http://interfacelift.com/wallpaper/downloads/date/widescreen_16:9/1920x1080/'
host = 'http://interfacelift.com'
saveDir = 'scraped'

# how many you want to download, 0 to download all
goal = 0

# go to your scrape url, note down how many pages there are. this should not exceed the max.
maxPages = 240

# hi guy from the future, if this script doesn't work, you probably need to update the regex.
tag = re.compile(r'<a href=\"(?P<url>.+)\"><img.+?src=\"/img_NEW/button_download')

if not os.path.exists(saveDir) :
    os.makedirs(saveDir)

# pointless, just makes the output a little nicer.
# formats goal if 0 to read 'about XX'
goalStr = str(goal)
if goal == 0:
    # round up maxPages*10 to next hundredth
    # shitty convoluted code, i know, but works
    goal = maxPages*10 + 99
    goal = goal - (goal % 100)
    goalStr = 'about %s' % str(goal)

done = 1
for page in range(1, maxPages+1):
    print '[+] opening page %s of %s' % (page, maxPages)

    url = scrapeUrl + "index" + str(page)

    # interfacelift returns a 403 forbidden unless you include a referer.
    headers = { 'User-agent' : "Mozilla/4.0 (compatible; MSIE 5.5; Windows NT)",
                'Referer': url}
    try:
        req = urllib2.Request(url, None, headers)
        f = urllib2.urlopen(req)

    except IOError, e:
        print 'Failed to open', url
        if hasattr(e, 'code'):
            print 'Error code:', e.code
        continue

    html = f.read()

    links = tag.finditer(html)

    for i in links:
        failed = False
        if goal == 0 or done <= goal :
            print '[-] scraping %d of %s' % (done, goalStr)
            dlurl = host + i.group('url')
            req = urllib2.Request(dlurl, None, headers)
            saveFile = os.path.join(saveDir, str(done)+'.jpg')
            with open(saveFile, 'wb') as f:
                try:
                    res = urllib2.urlopen(req)
                    f.write(res.read())
                except:
                    print 'failed'
                    failed = True
            if failed:
                try: os.remove(saveFile)
                except: pass

            done += 1
<|endoftext|>


<commit_before>from mongoengine import *
import pymongo

connect("test", host='127.0.0.1', port=27017, read_preference=pymongo.ReadPreference.PRIMARY)
default_document = """#!/bin/bash
sudo mkdir -p /opt/[[company]]
sudo chmod [[opt_level]] /opt/[[company]]
cd /opt/[[company]]
git checkout http://[[oauth]]@github.com/[[company]]/[[cookbooks]].git
echo "cookbook_path [\"/opt/[[company]]/[[cookbooks]]\"]" >> /opt/[[company]]/[[client_rb]]
[[chef_install]]
echo "Run your setup by running: chef-solo -j /opt/[[company]]/[[cookbooks]]/[[dev_setup_json]] -c /opt/[[company]]/[[client_rb]]"
"""

def get_defaults(class_to_check):
    rtn_dict = {}
    for k in class_to_check._db_field_map:
        try:
            rtn_dict[k] = getattr(class_to_check, k).default
        except AttributeError:
            rtn_dict[k] = None
        except Exception as e:
            print(e)
            rtn_dict[k] = None
    return rtn_dict

class Company(Document):
    name = StringField(required=True, unique=True)

class MainDoc(Document):
    company = ReferenceField(Company)
    user = LongField(required=True)

class DefaultDoc(Document):
    company = StringField(required=True)  # needs to resolve down to name later
    user = LongField(required=True)  # for filtering really
    oauth = StringField(required=True, unique=False)  # custom per setup, genreated at github.com
    opt_level = LongField(default=777)  # default 777
    cookbooks = StringField(default="cookbooks")  # name of the repo and checklout location: ie cookbooks
    dev_setup_json = StringField(default="dev_setup.json")  # requires extension: ie dev_setup.json
    client_rb = StringField(default="client.rb")  # default to client.rb
    chef_install = StringField(default="curl -L https://www.chef.io/chef/install.sh | sudo bash")  # for chef client
<commit_msg>changing the mongo host and ports to environment variables<commit_after>from mongoengine import *
import pymongo
import os

connect("miseenplace",
        host=os.getenv("MONGO_HOST", '127.0.0.1'),
        port=os.getenv("MONGO_PORT", 27017),
        read_preference=pymongo.ReadPreference.PRIMARY)
default_document = """#!/bin/bash
sudo mkdir -p /opt/[[company]]
sudo chmod [[opt_level]] /opt/[[company]]
cd /opt/[[company]]
git checkout http://[[oauth]]@github.com/[[company]]/[[cookbooks]].git
echo "cookbook_path [\"/opt/[[company]]/[[cookbooks]]\"]" >> /opt/[[company]]/[[client_rb]]
[[chef_install]]
echo "Run your setup by running: chef-solo -j /opt/[[company]]/[[cookbooks]]/[[dev_setup_json]] -c /opt/[[company]]/[[client_rb]]"
"""

def get_defaults(class_to_check):
    rtn_dict = {}
    for k in class_to_check._db_field_map:
        try:
            rtn_dict[k] = getattr(class_to_check, k).default
        except AttributeError:
            rtn_dict[k] = None
        except Exception as e:
            print(e)
            rtn_dict[k] = None
    return rtn_dict

class Company(Document):
    name = StringField(required=True, unique=True)

class MainDoc(Document):
    company = ReferenceField(Company)
    user = LongField(required=True)

class DefaultDoc(Document):
    company = StringField(required=True)  # needs to resolve down to name later
    user = LongField(required=True)  # for filtering really
    oauth = StringField(required=True, unique=False)  # custom per setup, genreated at github.com
    opt_level = LongField(default=777)  # default 777
    cookbooks = StringField(default="cookbooks")  # name of the repo and checklout location: ie cookbooks
    dev_setup_json = StringField(default="dev_setup.json")  # requires extension: ie dev_setup.json
    client_rb = StringField(default="client.rb")  # default to client.rb
    chef_install = StringField(default="curl -L https://www.chef.io/chef/install.sh | sudo bash")  # for chef client
<|endoftext|>


<commit_before>#!/bin/python
#    This file is part of recover-evtx.
#
#   Copyright 2013 Willi Ballenthin <william.ballenthin@mandiant.com>
#                    while at Mandiant <http://www.mandiant.com>
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
#   Version v0.1
import re
import logging

# TODO(wb): fallback to standard xml parser
from lxml import etree

def to_lxml(record):
    """
    @type record: Record
    """
    return etree.fromstring(
        "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"yes\" ?>%s" %
        (record.root().xml([]).encode("utf-8")))

def get_child(node, tag,
              ns="{http://schemas.microsoft.com/win/2004/08/events/event}"):
    """
    @type node: Element
    @type tag: str
    @type ns: str
    """
    return node.find("%s%s" % (ns, tag))

def get_eid(record):
    return get_child(get_child(to_lxml(record), "System"), "EventID").text

class Template(object):
    substitition_re = re.compile("\[(Conditional|Normal) Substitution\(index=(\d+), type=(\d+)\)\]")

    def __init__(self, eid, xml, offset):
        self._eid = eid
        self._xml = xml
        self._offset = offset

        self._cached_placeholders = None
        self._cached_id = None

    def get_xml(self):
        return self._xml

    def get_eid(self):
        return self._eid

    def get_id(self):
        if self._cached_id is not None:
            return self._cached_id
        ret = ["%s" % self._eid]
        for index, type_, mode in self._get_placeholders():
            if mode:
                mode_str = "c"
            else:
                mode_str = "n"
            ret.append("[%s|%s|%s]" % (index, type_, mode_str))
        self._cached_id = "-".join(ret)
        return self._cached_id

    def get_offset(self):
        return self._offset

    def _get_placeholders(self):
        """
        Get descriptors for each of the substitutions required by this
          template.

        Tuple schema: (index, type, is_conditional)

        @rtype: list of (int, int, boolean)
        """
        if self._cached_placeholders is not None:
            return self._cached_placeholders
        ret = []
        for mode, index, type_ in Template.substitition_re.findall(self._xml):
            ret.append((int(index), int(type_), mode == "Conditional"))
        self._cached_placeholders = sorted(ret, key=lambda p: p[0])
        return self._cached_placeholders

    def match_substitutions(self, substitutions):
        """
        Checks to see if the provided set of substitutions match the
          placeholder values required by this template.

        Note, this is only a best guess.  The number of substitutions
          *may* be greater than the number of available slots. So we
          must only check the slot and substitution types.

        Tuple schema: (index, type, value)

        @type substitutions: list of (int, int, str)
        @rtype: boolean
        """
        logger = logging.getLogger("match_substitutions")
        placeholders = self._get_placeholders()
        if len(placeholders) > len(substitutions):
            logger.debug("Failing on lens: %d vs %d",
                         len(placeholders), len(substitutions))
            return False
        if max(placeholders, key=lambda k: k[0])[0] > max(substitutions, key=lambda k: k[0])[0]:
            logger.debug("Failing on max index: %d vs %d",
                         max(placeholders, key=lambda k: k[0])[0],
                         max(substitutions, key=lambda k: k[0])[0])
            return False
        for index, type_, mode in placeholders:
            sub = substitutions[index]
            # substitutions should be sorted and index-able,
            #  but if not, fallback
            if sub[0] != index:
                for s in substitutions:
                    if s[0] == index:
                        sub = s
                        break
            if mode and sub[1] == 0:
                continue
            if sub[1] != type_:
                logger.debug("Failing on type comparison: %d vs %d (mode: %s)",
                             sub[1], type_, mode)
                return False
        return True

    escape_re = re.compile(r"\\\\(\d)")

    @staticmethod
    def _escape(value):
        """
        Escape the static value to be used in a regular expression
          subsititution. This processes any backreferences and
          makes them plain, escaped sequences.

        @type value: str
        @rtype: str
        """
        return Template.escape_re.sub(r"\\\\\\\\\1", re.escape(value))

    def insert_substitutions(self, substitutions):
        """
        Return a copy of the template with the given substitutions inserted.

        Implementation is a huge hack that depends on the
        brittle template_format() output.

        substitutions should be a list of (index:int, type:int, value:str)

        @type substitutions: list of (int, int, str)
        @rtype: str
        """
        ret = self._xml
        for index, type_, value in substitutions:
            from_pattern = "\[:?(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % index
            ret = re.sub(from_pattern, Template._escape(value), ret)
        return ret

    def serialize(self):
        ret = []
        ret.append("TEMPLATE")
        ret.append("ID: %s" % self.get_id())
        ret.append("EID: %s" % self.get_eid())
        ret.append("OFFSET: %s" % self.get_offset())
        ret.append("%s" % self.get_xml())
        return "\n".join(ret)

    @classmethod
    def deserialize(cls, text):
        for requirement in ["TEMPLATE", "EID: ", "ID: ", "<Event "]:
            if requirement not in text:
                return None
        template_lines = text.split("\n")
        #template_line = template_lines[0]
        id_line = template_lines[1]
        eid_line = template_lines[2]
        offset_line = template_lines[3]
        xml = "\n".join(template_lines[4:])
        _, __, id_ = id_line.partition(": ")
        _, __, eid = eid_line.partition(": ")
        _, __, offset = offset_line.partition(": ")
        eid = int(eid)
        offset = int(offset, 0x10)
        return cls(eid, xml, offset)

class TemplateEIDConflictError(Exception):
    def __init__(self, value):
        super(TemplateEIDConflictError, self).__init__(value)

class TemplateNotFoundError(Exception):
    def __init__(self, value):
        super(TemplateNotFoundError, self).__init__(value)

class TemplateDatabase(object):
    def __init__(self):
        # @type self._templates: {str: [Template]}
        # ID --> list matching Templates with the ID
        self._templates = {}
        # @type self._eid_map: {int: [str]}
        # EID --> list of ID
        self._eid_map = {}

    def add_template(self, template):
        """
        Merge a new Template into the database.

        @type template: template
        @rtype: None
        """
        id_ = template.get_id()
        if id_ in self._templates:
            found_existing = False
            for other in self._templates[id_]:
                if template.get_xml() == other.get_xml():
                    found_existing = True
                    break
            if not found_existing:
                self._templates[id_].append(template)
        else:
            self._templates[id_] = [template]
            if template.get_eid() in self._eid_map:
                self._eid_map[template.get_eid()].append(id_)
            else:
                self._eid_map[template.get_eid()] = [id_]

    def extend(self, other):
        """
        Merge another TemplateDatabase into this one.

        @type other: TemplateDatabase
        @rtype: None
        """
        for template_list in other._templates.values():
            for template in template_list:
                self.add_template(template)

    def get_template(self, eid, substitutions, exact_match=True):
        """
        Given an ID, attempt to pick the appropriate template.

        `exact_match` should only be unset during testing.
          Not for forensics.

        @type eid: int
        @type substitutions: list of (int, str)
        @rtype: Template
        @raises TemplateEIDConflictError
        @raises TemplateNotFoundError
        """
        if eid not in self._eid_map:
            raise TemplateNotFoundError(
                "No loaded templates with EID: %s" % eid)

        potential_templates = []
        for id_ in self._eid_map[eid]:
            potential_templates.extend(self._templates[id_])

        matching_templates = []
        for template in potential_templates:
            if template.match_substitutions(substitutions):
                matching_templates.append(template)

        if exact_match and len(matching_templates) > 1:
            matches = map(lambda t: t.get_id(), matching_templates)
            raise TemplateEIDConflictError("%d templates matched query for "
                                           "EID %d and substitutions: %s" % (len(matching_templates), eid, matches))
        if len(matching_templates) == 0:
            sig = str(eid) + "-" + "-".join(["[%d|%d| ]" % (i, j) for i, j in \
                                                 enumerate(map(lambda p: p[1], substitutions))])
            raise TemplateNotFoundError(
                "No loaded templates with given substitution signature: %s" % sig)

        return matching_templates[0]

    def serialize(self):
        """
        Create a string representation of this database, suitable
          for human review or deserialization.

        @rtype: str
        """
        ret = []
        for id_ in sorted(self._templates.keys()):
            for template in self._templates[id_]:
                ret.append(template.serialize())
                ret.append("\n\n")
        return "\n".join(ret)

    def deserialize(self, txt, warn_on_conflict=True):
        """
        Load a serialized TemplateDatabase into this one.

        Merges with whatever is currently in this database.

        @type txt: str
        @rtype: None
        @raises: TemplateEIDConflictError
        """
        self._templates = {}
        for template_txt in txt.split("TEMPLATE\n"):
            template = Template.deserialize("TEMPLATE\n" + template_txt.rstrip("\n"))
            if template is None:
                continue
            if template.get_id() in self._templates and warn_on_conflict:
                raise TemplateEIDConflictError("More than one template with ID %s" % template.get_id())
            elif template.get_id() in self._templates and not warn_on_conflict:
                self._templates[template.get_id()].append(template)
            else:
                self._templates[template.get_id()] = [template]
                if template.get_eid() in self._eid_map:
                    self._eid_map[template.get_eid()].append(template.get_id())
                else:
                    self._eid_map[template.get_eid()] = [template.get_id()]
<commit_msg>update regex for non-capturing group<commit_after>#!/bin/python
#    This file is part of recover-evtx.
#
#   Copyright 2013 Willi Ballenthin <william.ballenthin@mandiant.com>
#                    while at Mandiant <http://www.mandiant.com>
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
#   Version v0.1
import re
import logging

# TODO(wb): fallback to standard xml parser
from lxml import etree

def to_lxml(record):
    """
    @type record: Record
    """
    return etree.fromstring(
        "<?xml version=\"1.0\" encoding=\"utf-8\" standalone=\"yes\" ?>%s" %
        (record.root().xml([]).encode("utf-8")))

def get_child(node, tag,
              ns="{http://schemas.microsoft.com/win/2004/08/events/event}"):
    """
    @type node: Element
    @type tag: str
    @type ns: str
    """
    return node.find("%s%s" % (ns, tag))

def get_eid(record):
    return get_child(get_child(to_lxml(record), "System"), "EventID").text

class Template(object):
    substitition_re = re.compile("\[(Conditional|Normal) Substitution\(index=(\d+), type=(\d+)\)\]")

    def __init__(self, eid, xml, offset):
        self._eid = eid
        self._xml = xml
        self._offset = offset

        self._cached_placeholders = None
        self._cached_id = None

    def get_xml(self):
        return self._xml

    def get_eid(self):
        return self._eid

    def get_id(self):
        if self._cached_id is not None:
            return self._cached_id
        ret = ["%s" % self._eid]
        for index, type_, mode in self._get_placeholders():
            if mode:
                mode_str = "c"
            else:
                mode_str = "n"
            ret.append("[%s|%s|%s]" % (index, type_, mode_str))
        self._cached_id = "-".join(ret)
        return self._cached_id

    def get_offset(self):
        return self._offset

    def _get_placeholders(self):
        """
        Get descriptors for each of the substitutions required by this
          template.

        Tuple schema: (index, type, is_conditional)

        @rtype: list of (int, int, boolean)
        """
        if self._cached_placeholders is not None:
            return self._cached_placeholders
        ret = []
        for mode, index, type_ in Template.substitition_re.findall(self._xml):
            ret.append((int(index), int(type_), mode == "Conditional"))
        self._cached_placeholders = sorted(ret, key=lambda p: p[0])
        return self._cached_placeholders

    def match_substitutions(self, substitutions):
        """
        Checks to see if the provided set of substitutions match the
          placeholder values required by this template.

        Note, this is only a best guess.  The number of substitutions
          *may* be greater than the number of available slots. So we
          must only check the slot and substitution types.

        Tuple schema: (index, type, value)

        @type substitutions: list of (int, int, str)
        @rtype: boolean
        """
        logger = logging.getLogger("match_substitutions")
        placeholders = self._get_placeholders()
        if len(placeholders) > len(substitutions):
            logger.debug("Failing on lens: %d vs %d",
                         len(placeholders), len(substitutions))
            return False
        if max(placeholders, key=lambda k: k[0])[0] > max(substitutions, key=lambda k: k[0])[0]:
            logger.debug("Failing on max index: %d vs %d",
                         max(placeholders, key=lambda k: k[0])[0],
                         max(substitutions, key=lambda k: k[0])[0])
            return False
        for index, type_, mode in placeholders:
            sub = substitutions[index]
            # substitutions should be sorted and index-able,
            #  but if not, fallback
            if sub[0] != index:
                for s in substitutions:
                    if s[0] == index:
                        sub = s
                        break
            if mode and sub[1] == 0:
                continue
            if sub[1] != type_:
                logger.debug("Failing on type comparison: %d vs %d (mode: %s)",
                             sub[1], type_, mode)
                return False
        return True

    escape_re = re.compile(r"\\\\(\d)")

    @staticmethod
    def _escape(value):
        """
        Escape the static value to be used in a regular expression
          subsititution. This processes any backreferences and
          makes them plain, escaped sequences.

        @type value: str
        @rtype: str
        """
        return Template.escape_re.sub(r"\\\\\\\\\1", re.escape(value))

    def insert_substitutions(self, substitutions):
        """
        Return a copy of the template with the given substitutions inserted.

        Implementation is a huge hack that depends on the
        brittle template_format() output.

        substitutions should be a list of (index:int, type:int, value:str)

        @type substitutions: list of (int, int, str)
        @rtype: str
        """
        ret = self._xml
        for index, type_, value in substitutions:
            from_pattern = "\[(Normal|Conditional) Substitution\(index=%d, type=\d+\)\]" % index
            ret = re.sub(from_pattern, Template._escape(value), ret)
        return ret

    def serialize(self):
        ret = []
        ret.append("TEMPLATE")
        ret.append("ID: %s" % self.get_id())
        ret.append("EID: %s" % self.get_eid())
        ret.append("OFFSET: %s" % self.get_offset())
        ret.append("%s" % self.get_xml())
        return "\n".join(ret)

    @classmethod
    def deserialize(cls, text):
        for requirement in ["TEMPLATE", "EID: ", "ID: ", "<Event "]:
            if requirement not in text:
                return None
        template_lines = text.split("\n")
        #template_line = template_lines[0]
        id_line = template_lines[1]
        eid_line = template_lines[2]
        offset_line = template_lines[3]
        xml = "\n".join(template_lines[4:])
        _, __, id_ = id_line.partition(": ")
        _, __, eid = eid_line.partition(": ")
        _, __, offset = offset_line.partition(": ")
        eid = int(eid)
        offset = int(offset, 0x10)
        return cls(eid, xml, offset)

class TemplateEIDConflictError(Exception):
    def __init__(self, value):
        super(TemplateEIDConflictError, self).__init__(value)

class TemplateNotFoundError(Exception):
    def __init__(self, value):
        super(TemplateNotFoundError, self).__init__(value)

class TemplateDatabase(object):
    def __init__(self):
        # @type self._templates: {str: [Template]}
        # ID --> list matching Templates with the ID
        self._templates = {}
        # @type self._eid_map: {int: [str]}
        # EID --> list of ID
        self._eid_map = {}

    def add_template(self, template):
        """
        Merge a new Template into the database.

        @type template: template
        @rtype: None
        """
        id_ = template.get_id()
        if id_ in self._templates:
            found_existing = False
            for other in self._templates[id_]:
                if template.get_xml() == other.get_xml():
                    found_existing = True
                    break
            if not found_existing:
                self._templates[id_].append(template)
        else:
            self._templates[id_] = [template]
            if template.get_eid() in self._eid_map:
                self._eid_map[template.get_eid()].append(id_)
            else:
                self._eid_map[template.get_eid()] = [id_]

    def extend(self, other):
        """
        Merge another TemplateDatabase into this one.

        @type other: TemplateDatabase
        @rtype: None
        """
        for template_list in other._templates.values():
            for template in template_list:
                self.add_template(template)

    def get_template(self, eid, substitutions, exact_match=True):
        """
        Given an ID, attempt to pick the appropriate template.

        `exact_match` should only be unset during testing.
          Not for forensics.

        @type eid: int
        @type substitutions: list of (int, str)
        @rtype: Template
        @raises TemplateEIDConflictError
        @raises TemplateNotFoundError
        """
        if eid not in self._eid_map:
            raise TemplateNotFoundError(
                "No loaded templates with EID: %s" % eid)

        potential_templates = []
        for id_ in self._eid_map[eid]:
            potential_templates.extend(self._templates[id_])

        matching_templates = []
        for template in potential_templates:
            if template.match_substitutions(substitutions):
                matching_templates.append(template)

        if exact_match and len(matching_templates) > 1:
            matches = map(lambda t: t.get_id(), matching_templates)
            raise TemplateEIDConflictError("%d templates matched query for "
                                           "EID %d and substitutions: %s" % (len(matching_templates), eid, matches))
        if len(matching_templates) == 0:
            sig = str(eid) + "-" + "-".join(["[%d|%d| ]" % (i, j) for i, j in \
                                                 enumerate(map(lambda p: p[1], substitutions))])
            raise TemplateNotFoundError(
                "No loaded templates with given substitution signature: %s" % sig)

        return matching_templates[0]

    def serialize(self):
        """
        Create a string representation of this database, suitable
          for human review or deserialization.

        @rtype: str
        """
        ret = []
        for id_ in sorted(self._templates.keys()):
            for template in self._templates[id_]:
                ret.append(template.serialize())
                ret.append("\n\n")
        return "\n".join(ret)

    def deserialize(self, txt, warn_on_conflict=True):
        """
        Load a serialized TemplateDatabase into this one.

        Merges with whatever is currently in this database.

        @type txt: str
        @rtype: None
        @raises: TemplateEIDConflictError
        """
        self._templates = {}
        for template_txt in txt.split("TEMPLATE\n"):
            template = Template.deserialize("TEMPLATE\n" + template_txt.rstrip("\n"))
            if template is None:
                continue
            if template.get_id() in self._templates and warn_on_conflict:
                raise TemplateEIDConflictError("More than one template with ID %s" % template.get_id())
            elif template.get_id() in self._templates and not warn_on_conflict:
                self._templates[template.get_id()].append(template)
            else:
                self._templates[template.get_id()] = [template]
                if template.get_eid() in self._eid_map:
                    self._eid_map[template.get_eid()].append(template.get_id())
                else:
                    self._eid_map[template.get_eid()] = [template.get_id()]
<|endoftext|>


<commit_before>from sqlalchemy import Column, Date, Integer, Float, String
from sqlalchemy import ForeignKeyConstraint, ForeignKey, Index, Text
from sqlalchemy.orm import deferred, relationship
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.ext.declarative import declarative_base

# Extend the Base >>>>>>
Base = declarative_base()

def init(self, *args, **kwargs):
    for i, arg in enumerate(args):
        self.__dict__[self.kw[i]] = arg
    for k, v in kwargs.iteritems():
        self.__dict__[k] = v
Base.__init__ = init
# <<<<<<

# PATENT ---------------------------

class Patent(Base):
    __tablename__ = "patent"
    uuid = Column(Integer, primary_key=True)
    type = Column(String(20))
    number = Column(String(20))
    country = Column(String(20))
    date = Column(Date)
    abstract = deferred(Column(Text))
    title = deferred(Column(Text))
    kind = Column(String(10))
    claims = Column(Integer)

    application = relationship("Application", uselist=False, backref="patent")
    classes = relationship("USPC", backref="patent")

    assignees = relationship("Assignee", backref="patent")
    inventors = relationship("Inventor", backref="patent")
    lawyers = relationship("Lawyer", backref="patent")

    citations = relationship(
        "Citation",
        primaryjoin="Patent.uuid == Citation.patent_uuid",
        backref="patent")
    citedby = relationship(
        "Citation",
        primaryjoin="Patent.uuid == Citation.citation_uuid",
        backref="citation")
    otherreferences = relationship("OtherReference", backref="patent")
    usreldocs = relationship("USRelDoc", backref="patent")

    kw = ["type", "number", "country", "date",
          "abstract", "title", "kind", "claims"]
    __table_args__ = (
        Index("pat_idx1", "type", "number", unique=True),
        Index("pat_idx2", "date"),
    )

class Application(Base):
    __tablename__ = "application"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    type = Column(String(20))
    number = Column(String(20))
    country = Column(String(20))
    date = Column(Date)
    kw = ["type", "number", "country", "date"]
    __table_args__ = (
        Index("app_idx1", "type", "number", unique=True),
        Index("app_idx2", "date"),
    )
#Index('pat_idx3', Patent.date__grant)

# SUPPORT --------------------------

class Location(Base):
    #TODO: Anyway we can consolidate to 1 PrimaryKey
    __tablename__ = "location"
    city = Column(String(128), primary_key=True)
    state = Column(String(10), index=True, primary_key=True)
    country = Column(String(10), index=True, primary_key=True)
    latitude = Column(Float)
    longitude = Column(Float)
    inventors = relationship("Inventor", backref="location")
    assignees = relationship("Assignee", backref="location")
    kw = ["city", "state", "country", "longitude", "latitude"]
    __table_args__ = (
        Index("loc_idx1", "latitude", "longitude"),
    )

    @hybrid_property
    def address(self):
        addy = []
        if self.city:
            addy.append(self.city)
        if self.state:
            addy.append(self.state)
        if self.country:
            addy.append(self.country)
        return ", ".join(addy)

# OBJECTS --------------------------

class Assignee(Base):
    __tablename__ = "assignee"
    uuid = Column(Integer, primary_key=True)
    #location_uuid = Column(Integer, ForeignKey("location.uuid"))
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    type = Column(String(10))
    organization = Column(String(256))
    name_first = Column(String(64))
    name_last = Column(String(64))
    location_city = Column(String(128))
    location_state = Column(String(10), index=True)
    location_country = Column(String(10), index=True)
    sequence = Column(Integer, index=True)
    kw = ["sequence"]
    __table_args__ = (
        ForeignKeyConstraint(
            [loc_city, loc_state, loc_country],
            [Location.city, Location.state, Location.country]
        ),
    )

    def asg(self, *args):
        self.organization = args[0]
        self.type = args[1]

    def person(self, *args):
        self.name_last = args[0]
        self.name_first = args[1]

class Inventor(Base):
    __tablename__ = "inventor"
    uuid = Column(Integer, primary_key=True)
    #location_uuid = Column(Integer, ForeignKey("location.uuid"))
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    name_last = Column(String(64))
    name_first = Column(String(64))
    location_city = Column(String(128))
    location_state = Column(String(10), index=True)
    location_country = Column(String(10), index=True)
    sequence = Column(Integer, index=True)
    kw = ["sequence", "name_last", "name_first", "nationality"]
    __table_args__ = (
        ForeignKeyConstraint(
            [loc_city, loc_state, loc_country],
            [Location.city, Location.state, Location.country]
        ),
    )

    @hybrid_property
    def name_full(self):
        return "{first} {last}".format(
            first=self.name_first,
            last=self.name_last)

class Lawyer(Base):
    __tablename__ = "lawyer"
    uuid = Column(Integer, primary_key=True)
    #location_uuid = Column(Integer, ForeignKey("location.uuid"))
    id = Column(String(64), index=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    name_last = Column(String(64))
    name_first = Column(String(64))
    organization = Column(String(64))
    country = Column(String(10))
    sequence = Column(Integer, index=True)
    kw = ["sequence", "name_last", "name_first", "organization", "country"]

    @hybrid_property
    def name_full(self):
        return "{first} {last}".format(
            first=self.name_first,
            last=self.name_last)

# REFERENCES -----------------------

class Citation(Base):
    """
    Two types of citations?
    """
    __tablename__ = "citation"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    citation_uuid = Column(Integer, ForeignKey("patent.uuid"))
    date = Column(Date, index=True)
    number = Column(String(20), index=True)
    name = Column(String(64))
    kind = Column(String(10))
    country = Column(String(10), index=True)
    category = Column(String(20), index=True)
    sequence = Column(Integer, index=True)
    kw = ["sequence", "date", "number", "name",
          "kind", "country", "category"]

class OtherReference(Base):
    __tablename__ = "otherreference"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    text = Column(Text)
    sequence = Column(Integer, index=True)
    kw = ["sequence", "text"]

class USRelDoc(Base):
    __tablename__ = "usreldoc"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))

# CLASSIFICATIONS ------------------

class USPC(Base):
    __tablename__ = "uspc"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    mainclass_id = Column(String(10), ForeignKey("mainclass.id"))
    subclass_id = Column(String(10), ForeignKey("subclass.id"))
    sequence = Column(Integer, index=True)
    kw = ["sequence"]

class MainClass(Base):
    __tablename__ = "mainclass"
    id = Column(String(20), primary_key=True)
    title = Column(String(256))
    text = Column(String(256))
    uspc = relationship("USPC", backref="mainclass")
    kw = ["id", "title", "text"]

class SubClass(Base):
    __tablename__ = "subclass"
    id = Column(String(20), primary_key=True)
    title = Column(String(256))
    text = Column(String(256))
    uspc = relationship("USPC", backref="subclass")
    kw = ["id", "title", "text"]
<commit_msg>remodified schema.py<commit_after>from sqlalchemy import Column, Date, Integer, Float, String
from sqlalchemy import ForeignKeyConstraint, ForeignKey, Index, Text
from sqlalchemy.orm import deferred, relationship
from sqlalchemy.ext.hybrid import hybrid_property
from sqlalchemy.ext.declarative import declarative_base

# Extend the Base >>>>>>
Base = declarative_base()

def init(self, *args, **kwargs):
    for i, arg in enumerate(args):
        self.__dict__[self.kw[i]] = arg
    for k, v in kwargs.iteritems():
        self.__dict__[k] = v
Base.__init__ = init
# <<<<<<

# PATENT ---------------------------

class Patent(Base):
    __tablename__ = "patent"
    uuid = Column(Integer, primary_key=True)
    type = Column(String(20))
    number = Column(String(20))
    country = Column(String(20))
    date = Column(Date)
    abstract = deferred(Column(Text))
    title = deferred(Column(Text))
    kind = Column(String(10))
    claims = Column(Integer)

    application = relationship("Application", uselist=False, backref="patent")
    classes = relationship("USPC", backref="patent")

    assignees = relationship("Assignee", backref="patent")
    inventors = relationship("Inventor", backref="patent")
    lawyers = relationship("Lawyer", backref="patent")

    citations = relationship(
        "Citation",
        primaryjoin="Patent.uuid == Citation.patent_uuid",
        backref="patent")
    citedby = relationship(
        "Citation",
        primaryjoin="Patent.uuid == Citation.citation_uuid",
        backref="citation")
    otherreferences = relationship("OtherReference", backref="patent")
    usreldocs = relationship("USRelDoc", backref="patent")

    kw = ["type", "number", "country", "date",
          "abstract", "title", "kind", "claims"]
    __table_args__ = (
        Index("pat_idx1", "type", "number", unique=True),
        Index("pat_idx2", "date"),
    )

class Application(Base):
    __tablename__ = "application"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    type = Column(String(20))
    number = Column(String(20))
    country = Column(String(20))
    date = Column(Date)
    kw = ["type", "number", "country", "date"]
    __table_args__ = (
        Index("app_idx1", "type", "number", unique=True),
        Index("app_idx2", "date"),
    )
#Index('pat_idx3', Patent.date__grant)

# SUPPORT --------------------------

class Location(Base):
    #TODO: Anyway we can consolidate to 1 PrimaryKey
    __tablename__ = "location"
    city = Column(String(128), primary_key=True)
    state = Column(String(10), index=True, primary_key=True)
    country = Column(String(10), index=True, primary_key=True)
    latitude = Column(Float)
    longitude = Column(Float)
    inventors = relationship("Inventor", backref="location")
    assignees = relationship("Assignee", backref="location")
    kw = ["city", "state", "country", "longitude", "latitude"]
    __table_args__ = (
        Index("loc_idx1", "latitude", "longitude"),
    )

    @hybrid_property
    def address(self):
        addy = []
        if self.city:
            addy.append(self.city)
        if self.state:
            addy.append(self.state)
        if self.country:
            addy.append(self.country)
        return ", ".join(addy)

# OBJECTS --------------------------

class Assignee(Base):
    __tablename__ = "assignee"
    uuid = Column(Integer, primary_key=True)
    #location_uuid = Column(Integer, ForeignKey("location.uuid"))
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    type = Column(String(10))
    organization = Column(String(256))
    name_first = Column(String(64))
    name_last = Column(String(64))
    location_city = Column(String(128))
    location_state = Column(String(10), index=True)
    location_country = Column(String(10), index=True)
    sequence = Column(Integer, index=True)
    kw = ["sequence"]
    __table_args__ = (
        ForeignKeyConstraint(
            [location_city, location_state, location_country],
            [Location.city, Location.state, Location.country]
        ),
    )

    def asg(self, *args):
        self.organization = args[0]
        self.type = args[1]

    def person(self, *args):
        self.name_last = args[0]
        self.name_first = args[1]

class Inventor(Base):
    __tablename__ = "inventor"
    uuid = Column(Integer, primary_key=True)
    #location_uuid = Column(Integer, ForeignKey("location.uuid"))
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    name_last = Column(String(64))
    name_first = Column(String(64))
    location_city = Column(String(128))
    location_state = Column(String(10), index=True)
    location_country = Column(String(10), index=True)
    sequence = Column(Integer, index=True)
    kw = ["sequence", "name_last", "name_first", "nationality"]
    __table_args__ = (
        ForeignKeyConstraint(
            [location_city, location_state, location_country],
            [Location.city, Location.state, Location.country]
        ),
    )

    @hybrid_property
    def name_full(self):
        return "{first} {last}".format(
            first=self.name_first,
            last=self.name_last)

class Lawyer(Base):
    __tablename__ = "lawyer"
    uuid = Column(Integer, primary_key=True)
    #location_uuid = Column(Integer, ForeignKey("location.uuid"))
    id = Column(String(64), index=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    name_last = Column(String(64))
    name_first = Column(String(64))
    organization = Column(String(64))
    country = Column(String(10))
    sequence = Column(Integer, index=True)
    kw = ["sequence", "name_last", "name_first", "organization", "country"]

    @hybrid_property
    def name_full(self):
        return "{first} {last}".format(
            first=self.name_first,
            last=self.name_last)

# REFERENCES -----------------------

class Citation(Base):
    """
    Two types of citations?
    """
    __tablename__ = "citation"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    citation_uuid = Column(Integer, ForeignKey("patent.uuid"))
    date = Column(Date, index=True)
    number = Column(String(20), index=True)
    name = Column(String(64))
    kind = Column(String(10))
    country = Column(String(10), index=True)
    category = Column(String(20), index=True)
    sequence = Column(Integer, index=True)
    kw = ["sequence", "date", "number", "name",
          "kind", "country", "category"]

class OtherReference(Base):
    __tablename__ = "otherreference"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    text = Column(Text)
    sequence = Column(Integer, index=True)
    kw = ["sequence", "text"]

class USRelDoc(Base):
    __tablename__ = "usreldoc"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))

# CLASSIFICATIONS ------------------

class USPC(Base):
    __tablename__ = "uspc"
    uuid = Column(Integer, primary_key=True)
    patent_uuid = Column(Integer, ForeignKey("patent.uuid"))
    mainclass_id = Column(String(10), ForeignKey("mainclass.id"))
    subclass_id = Column(String(10), ForeignKey("subclass.id"))
    sequence = Column(Integer, index=True)
    kw = ["sequence"]

class MainClass(Base):
    __tablename__ = "mainclass"
    id = Column(String(20), primary_key=True)
    title = Column(String(256))
    text = Column(String(256))
    uspc = relationship("USPC", backref="mainclass")
    kw = ["id", "title", "text"]

class SubClass(Base):
    __tablename__ = "subclass"
    id = Column(String(20), primary_key=True)
    title = Column(String(256))
    text = Column(String(256))
    uspc = relationship("USPC", backref="subclass")
    kw = ["id", "title", "text"]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# This file is part of Invenio.
# Copyright (C) 2017 CERN.
#
# Invenio is free software; you can redistribute it
# and/or modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 2 of the
# License, or (at your option) any later version.
#
# Invenio is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Invenio; if not, write to the
# Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston,
# MA 02111-1307, USA.
#
# In applying this license, CERN does not
# waive the privileges and immunities granted to it by virtue of its status
# as an Intergovernmental Organization or submit itself to any jurisdiction.

"""Query processing classes."""

from datetime import datetime

import dateutil.parser
import six
from elasticsearch_dsl import Search
from invenio_search import current_search_client

from .errors import InvalidRequestInputError

class ESQuery(object):
    """Elasticsearch query."""

    def __init__(self, query_name, doc_type, index, client=None,
                 *args, **kwargs):
        """Constructor.

        :param doc_type: queried document type.
        :param index: queried index.
        :param client: elasticsearch client used to query.
        """
        super(ESQuery, self).__init__()
        self.index = index
        self.client = client or current_search_client
        self.query_name = query_name
        self.doc_type = doc_type

    def extract_date(self, date):
    """Extract date from string if necessary.

    :returns: the extracted date.
    """
    if isinstance(date, six.string_types):
        try:
            date = dateutil.parser.parse(date)
        except ValueError:
            raise ValueError(
                'Invalid date format for statistic {}.'
            ).format(self.query_name)
    if not isinstance(date, datetime):
        raise TypeError(
            'Invalid date type for statistic {}.'
        ).format(self.query_name)
    return date

    def run(self, *args, **kwargs):
        """Run the query."""
        raise NotImplementedError()

class ESDateHistogramQuery(ESQuery):
    """Elasticsearch date histogram query."""

    allowed_intervals = ['year', 'quarter', 'month', 'week', 'day']
    """Allowed intervals for the histogram aggregation."""

    def __init__(self, time_field='timestamp', copy_fields=None,
                 query_modifiers=None, required_filters=None,
                 metric_fields=None, *args, **kwargs):
        """Constructor.

        :param time_field: name of the timestamp field.
        :param copy_fields: list of fields to copy from the top hit document
            into the resulting aggregation item.
        :param query_modifiers: List of functions accepting a ``query`` and
            ``**kwargs`` (same as provided to the ``run`` method), that will
            be applied to the aggregation query.
        :param required_filters: Dict of "mandatory query parameter" ->
            "filtered field".
        :param metric_fields: Dict of "destination field" ->
            tuple("metric type", "source field").
        """
        super(ESDateHistogramQuery, self).__init__(*args, **kwargs)
        self.time_field = time_field
        self.copy_fields = copy_fields or {}
        self.query_modifiers = query_modifiers or []
        self.required_filters = required_filters or {}
        self.metric_fields = metric_fields or {}

    def validate_arguments(self, interval, start_date, end_date, **kwargs):
        """Validate query arguments."""
        if interval not in self.allowed_intervals:
            raise InvalidRequestInputError(
                'Invalid aggregation time interval for statistic {}.'
            ).format(self.query_name)
        if set(kwargs) < set(self.required_filters):
            raise InvalidRequestInputError(
                'Missing one of the required parameters {0} in '
                'query {1}'.format(set(self.required_filters.keys()),
                                   self.query_name)
            )

    def build_query(self, interval, start_date, end_date, **kwargs):
        """Build the elasticsearch query."""
        agg_query = Search(using=self.client,
                           index=self.index,
                           doc_type=self.doc_type)[0:0]
        if start_date is not None or end_date is not None:
            time_range = {}
            if start_date is not None:
                time_range['gte'] = start_date.isoformat()
            if end_date is not None:
                time_range['lte'] = end_date.isoformat()
            agg_query = agg_query.filter(
                'range',
                **{self.time_field: time_range})

        for modifier in self.query_modifiers:
            self.agg_query = modifier(self.agg_query, **kwargs)

            'histogram',
            'date_histogram',
            field=self.time_field,
            interval=interval
        )
        hist_agg.metric('total', 'sum', field='count')

        for destination, (metric, field) in self.metric_fields.items():
            hist_agg.metric(destination, metric, field=field)

        if self.copy_fields:
            hist_agg.metric(
                'top_hit', 'top_hits', size=1, sort={'timestamp': 'desc'}
            )

        for query_param, filtered_field in self.required_filters.items():
            if query_param in kwargs:
                agg_query = agg_query.filter(
                    'term', **{filtered_field: kwargs[query_param]}
                )

        return agg_query

    def process_query_result(self, result, interval, start_date, end_date):
        """Build the result using the query result."""
        def build_bucket(agg):
            result = dict(
                value=agg['total']['value'],
                key=agg['key'],
            )
            if self.copy_fields and agg['top_hit']['hits']['hits']:
                doc = agg['top_hit']['hits']['hits'][0]['_source']
                for destination, source in self.copy_fields.items():
                    if isinstance(source, six.string_types):
                        result[destination] = doc[source]
                    else:
                        result[destination] = source(
                            result,
                            doc
                        )
            if self.metric_fields:
                for destination in self.metric_fields:
                    result[destination] = agg[destination]['value']

            return result
        return dict(
            type='bucket',
            key_type='date',
            interval=interval,
            buckets=list(map(build_bucket,
                             result['aggregations']['histogram']['buckets']))
        )

    def run(self, interval='day', start_date=None,
            end_date=None, **kwargs):
        """Run the query."""
        start_date = self.extract_date(start_date) if start_date else None
        end_date = self.extract_date(end_date) if end_date else None
        self.validate_arguments(interval, start_date, end_date, **kwargs)

        agg_query = self.build_query(interval, start_date,
                                     end_date, **kwargs)
        query_result = agg_query.execute().to_dict()
        res = self.process_query_result(query_result, interval,
                                        start_date, end_date)
        return res

class ESTermsQuery(ESQuery):
    """Elasticsearch sum query."""

    def __init__(self, time_field='timestamp', copy_fields=None,
                 query_modifiers=None, required_filters=None,
                 aggregated_fields=None, metric_fields=None, *args, **kwargs):
        """Constructor.

        :param time_field: name of the timestamp field.
        :param copy_fields: list of fields to copy from the top hit document
            into the resulting aggregation item.
        :param query_modifiers: List of functions accepting a ``query`` and
            ``**kwargs`` (same as provided to the ``run`` method), that will
            be applied to the aggregation query.
        :param required_filters: Dict of "mandatory query parameter" ->
            "filtered field".
        :param aggregated_fields: List of fields which will be used in the
            terms aggregations.
        :param metric_fields: Dict of "destination field" ->
            tuple("metric type", "source field").
        """
        super(ESTermsQuery, self).__init__(*args, **kwargs)
        self.time_field = time_field
        self.copy_fields = copy_fields or {}
        self.query_modifiers = query_modifiers or []
        self.required_filters = required_filters or {}
        self.aggregated_fields = aggregated_fields or []
        self.metric_fields = metric_fields or {}
        assert len(self.aggregated_fields) > 0

    def validate_arguments(self, start_date, end_date, **kwargs):
        """Validate query arguments."""
        if set(kwargs) < set(self.required_filters):
            raise InvalidRequestInputError(
                'Missing one of the required parameters {0} in '
                'query {1}'.format(set(self.required_filters.keys()),
                                   self.query_name)
            )

    def build_query(self, start_date, end_date, **kwargs):
        """Build the elasticsearch query."""
        agg_query = Search(using=self.client,
                           index=self.index,
                           doc_type=self.doc_type)[0:0]
        if start_date is not None or end_date is not None:
            time_range = {}
            if start_date is not None:
                time_range['gte'] = start_date.isoformat()
            if end_date is not None:
                time_range['lte'] = end_date.isoformat()
            agg_query = agg_query.filter(
                'range',
                **{self.time_field: time_range})

        for modifier in self.query_modifiers:
            self.agg_query = modifier(self.agg_query, **kwargs)

        for destination, (metric, field) in self.metric_fields.items():
            term_agg.metric(destination, metric, field=field)

        if self.copy_fields:
            term_agg.metric(
                'top_hit', 'top_hits', size=1, sort={'timestamp': 'desc'}
            )

        for query_param, filtered_field in self.required_filters.items():
            if query_param in kwargs:
                agg_query = agg_query.filter(
                    'term', **{filtered_field: kwargs[query_param]}
                )

        return agg_query

    def process_query_result(self, result, start_date, end_date):
        """Build the result using the query result."""
        def build_buckets(agg, fields, res):
            """Build recursively result buckets."""
            if fields:
                field = fields[0]
                res.update(
                    type='bucket',
                    field=field,
                    key_type='terms',
                    buckets=list(map(
                        lambda sub: build_buckets(sub, fields[1:],
                                                  dict(key=sub['key'])),
                        agg[field]['buckets']))
                )
            else:
                res.update(
                    value=agg['total']['value'],
                    key=agg['key'],
                )
                if self.copy_fields and agg['top_hit']['hits']['hits']:
                    doc = agg['top_hit']['hits']['hits'][0]['_source']
                    for destination, source in self.copy_fields.items():
                        if isinstance(source, six.string_types):
                            res[destination] = doc[source]
                        else:
                            res[destination] = source(
                                res,
                                doc
                            )
                if self.metric_fields:
                    for destination in self.metric_fields:
                        result[destination] = agg[destination]['value']
            return res

        return build_buckets(result['aggregations'], self.aggregated_fields,
                             dict())

    def run(self, start_date=None, end_date=None, **kwargs):
        """Run the query."""
        start_date = self.extract_date(start_date) if start_date else None
        end_date = self.extract_date(end_date) if end_date else None
        self.validate_arguments(start_date, end_date, **kwargs)

        agg_query = self.build_query(start_date, end_date, **kwargs)
        query_result = agg_query.execute().to_dict()
        res = self.process_query_result(query_result, start_date, end_date)
        return res
<commit_msg>queries: refactor metric aggregations<commit_after># -*- coding: utf-8 -*-
#
# This file is part of Invenio.
# Copyright (C) 2017 CERN.
#
# Invenio is free software; you can redistribute it
# and/or modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 2 of the
# License, or (at your option) any later version.
#
# Invenio is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Invenio; if not, write to the
# Free Software Foundation, Inc., 59 Temple Place, Suite 330, Boston,
# MA 02111-1307, USA.
#
# In applying this license, CERN does not
# waive the privileges and immunities granted to it by virtue of its status
# as an Intergovernmental Organization or submit itself to any jurisdiction.

"""Query processing classes."""

from datetime import datetime

import dateutil.parser
import six
from elasticsearch_dsl import Search
from invenio_search import current_search_client

from .errors import InvalidRequestInputError

class ESQuery(object):
    """Elasticsearch query."""

    def __init__(self, query_name, doc_type, index, client=None,
                 *args, **kwargs):
        """Constructor.

        :param doc_type: queried document type.
        :param index: queried index.
        :param client: elasticsearch client used to query.
        """
        super(ESQuery, self).__init__()
        self.index = index
        self.client = client or current_search_client
        self.query_name = query_name
        self.doc_type = doc_type

    def extract_date(self, date):
        """Extract date from string if necessary.

        :returns: the extracted date.
        """
        if isinstance(date, six.string_types):
            try:
                date = dateutil.parser.parse(date)
            except ValueError:
                raise ValueError(
                    'Invalid date format for statistic {}.'
                ).format(self.query_name)
        if not isinstance(date, datetime):
            raise TypeError(
                'Invalid date type for statistic {}.'
            ).format(self.query_name)
        return date

    def run(self, *args, **kwargs):
        """Run the query."""
        raise NotImplementedError()

class ESDateHistogramQuery(ESQuery):
    """Elasticsearch date histogram query."""

    allowed_intervals = ['year', 'quarter', 'month', 'week', 'day']
    """Allowed intervals for the histogram aggregation."""

    def __init__(self, time_field='timestamp', copy_fields=None,
                 query_modifiers=None, required_filters=None,
                 metric_fields=None, *args, **kwargs):
        """Constructor.

        :param time_field: name of the timestamp field.
        :param copy_fields: list of fields to copy from the top hit document
            into the resulting aggregation.
        :param query_modifiers: List of functions accepting a ``query`` and
            ``**kwargs`` (same as provided to the ``run`` method), that will
            be applied to the aggregation query.
        :param required_filters: Dict of "mandatory query parameter" ->
            "filtered field".
        :param metric_fields: Dict of "destination field" ->
            tuple("metric type", "source field").
        """
        super(ESDateHistogramQuery, self).__init__(*args, **kwargs)
        self.time_field = time_field
        self.copy_fields = copy_fields or {}
        self.query_modifiers = query_modifiers or []
        self.required_filters = required_filters or {}
        self.metric_fields = metric_fields or {'value': ('sum', 'count')}

    def validate_arguments(self, interval, start_date, end_date, **kwargs):
        """Validate query arguments."""
        if interval not in self.allowed_intervals:
            raise InvalidRequestInputError(
                'Invalid aggregation time interval for statistic {}.'
            ).format(self.query_name)
        if set(kwargs) < set(self.required_filters):
            raise InvalidRequestInputError(
                'Missing one of the required parameters {0} in '
                'query {1}'.format(set(self.required_filters.keys()),
                                   self.query_name)
            )

    def build_query(self, interval, start_date, end_date, **kwargs):
        """Build the elasticsearch query."""
        agg_query = Search(using=self.client,
                           index=self.index,
                           doc_type=self.doc_type)[0:0]
        if start_date is not None or end_date is not None:
            time_range = {}
            if start_date is not None:
                time_range['gte'] = start_date.isoformat()
            if end_date is not None:
                time_range['lte'] = end_date.isoformat()
            agg_query = agg_query.filter(
                'range',
                **{self.time_field: time_range})

        for modifier in self.query_modifiers:
            self.agg_query = modifier(self.agg_query, **kwargs)

        base_agg = agg_query.aggs.bucket(
            'histogram',
            'date_histogram',
            field=self.time_field,
            interval=interval
        )

        for destination, (metric, field) in self.metric_fields.items():
            base_agg.metric(destination, metric, field=field)

        if self.copy_fields:
            base_agg.metric(
                'top_hit', 'top_hits', size=1, sort={'timestamp': 'desc'}
            )

        for query_param, filtered_field in self.required_filters.items():
            if query_param in kwargs:
                agg_query = agg_query.filter(
                    'term', **{filtered_field: kwargs[query_param]}
                )

        return agg_query

    def process_query_result(self, query_result, interval,
                             start_date, end_date):
        """Build the result using the query result."""
        def build_buckets(agg):
            """Build recursively result buckets."""
            bucket_result = dict(
                key=agg['key'],
                date=agg['key_as_string'],
            )
            for metric in self.metric_fields:
                bucket_result[metric] = agg[metric]['value']
            if self.copy_fields and agg['top_hit']['hits']['hits']:
                doc = agg['top_hit']['hits']['hits'][0]['_source']
                for destination, source in self.copy_fields.items():
                    if isinstance(source, six.string_types):
                        bucket_result[destination] = doc[source]
                    else:
                        bucket_result[destination] = source(bucket_result, doc)
            return bucket_result

        # Add copy_fields
        buckets = query_result['aggregations']['histogram']['buckets']
        return dict(
            interval=interval,
            key_type='date',
            start_date=start_date.isoformat(),
            end_date=end_date.isoformat(),
            buckets=[build_buckets(b) for b in buckets]
        )

    def run(self, interval='day', start_date=None,
            end_date=None, **kwargs):
        """Run the query."""
        start_date = self.extract_date(start_date) if start_date else None
        end_date = self.extract_date(end_date) if end_date else None
        self.validate_arguments(interval, start_date, end_date, **kwargs)

        agg_query = self.build_query(interval, start_date,
                                     end_date, **kwargs)
        query_result = agg_query.execute().to_dict()
        res = self.process_query_result(query_result, interval,
                                        start_date, end_date)
        return res

class ESTermsQuery(ESQuery):
    """Elasticsearch sum query."""

    def __init__(self, time_field='timestamp', copy_fields=None,
                 query_modifiers=None, required_filters=None,
                 aggregated_fields=None, metric_fields=None, *args, **kwargs):
        """Constructor.

        :param time_field: name of the timestamp field.
        :param copy_fields: list of fields to copy from the top hit document
            into the resulting aggregation.
        :param query_modifiers: List of functions accepting a ``query`` and
            ``**kwargs`` (same as provided to the ``run`` method), that will
            be applied to the aggregation query.
        :param required_filters: Dict of "mandatory query parameter" ->
            "filtered field".
        :param aggregated_fields: List of fields which will be used in the
            terms aggregations.
        :param metric_fields: Dict of "destination field" ->
            tuple("metric type", "source field").
        """
        super(ESTermsQuery, self).__init__(*args, **kwargs)
        self.time_field = time_field
        self.copy_fields = copy_fields or {}
        self.query_modifiers = query_modifiers or []
        self.required_filters = required_filters or {}
        self.aggregated_fields = aggregated_fields or []
        self.metric_fields = metric_fields or {'value': ('sum', 'count')}

    def validate_arguments(self, start_date, end_date, **kwargs):
        """Validate query arguments."""
        if set(kwargs) < set(self.required_filters):
            raise InvalidRequestInputError(
                'Missing one of the required parameters {0} in '
                'query {1}'.format(set(self.required_filters.keys()),
                                   self.query_name)
            )

    def build_query(self, start_date, end_date, **kwargs):
        """Build the elasticsearch query."""
        agg_query = Search(using=self.client,
                           index=self.index,
                           doc_type=self.doc_type)[0:0]
        if start_date is not None or end_date is not None:
            time_range = {}
            if start_date is not None:
                time_range['gte'] = start_date.isoformat()
            if end_date is not None:
                time_range['lte'] = end_date.isoformat()
            agg_query = agg_query.filter(
                'range',
                **{self.time_field: time_range})

        for modifier in self.query_modifiers:
            self.agg_query = modifier(self.agg_query, **kwargs)

        base_agg = agg_query.aggs

        def _apply_metric_aggs(agg):
            for destination, (metric, field) in self.metric_fields.items():
                agg.metric(destination, metric, field=field)

        _apply_metric_aggs(base_agg)
        if self.aggregated_fields:
            cur_agg = base_agg
            for term in self.aggregated_fields:
                cur_agg = cur_agg.bucket(term, 'terms', field=term, size=0)
                _apply_metric_aggs(cur_agg)

        if self.copy_fields:
            base_agg.metric(
                'top_hit', 'top_hits', size=1, sort={'timestamp': 'desc'}
            )

        for query_param, filtered_field in self.required_filters.items():
            if query_param in kwargs:
                agg_query = agg_query.filter(
                    'term', **{filtered_field: kwargs[query_param]}
                )

        return agg_query

    def process_query_result(self, query_result, start_date, end_date):
        """Build the result using the query result."""
        def build_buckets(agg, fields, bucket_result):
            """Build recursively result buckets."""
            # Add metric results for current bucket
            for metric in self.metric_fields:
                bucket_result[metric] = agg[metric]['value']
            if fields:
                current_level = fields[0]
                bucket_result.update(dict(
                    type='bucket',
                    field=current_level,
                    key_type='terms',
                    buckets=[build_buckets(b, fields[1:], dict(key=b['key']))
                             for b in agg[current_level]['buckets']]
                ))
            return bucket_result

        # Add copy_fields
        aggs = query_result['aggregations']
        result = dict(
            start_date=start_date.isoformat(),
            end_date=end_date.isoformat(),
        )
        if self.copy_fields and aggs['top_hit']['hits']['hits']:
            doc = aggs['top_hit']['hits']['hits'][0]['_source']
            for destination, source in self.copy_fields.items():
                if isinstance(source, six.string_types):
                    result[destination] = doc[source]
                else:
                    result[destination] = source(result, doc)

        return build_buckets(aggs, self.aggregated_fields, result)

    def run(self, start_date=None, end_date=None, **kwargs):
        """Run the query."""
        start_date = self.extract_date(start_date) if start_date else None
        end_date = self.extract_date(end_date) if end_date else None
        self.validate_arguments(start_date, end_date, **kwargs)

        agg_query = self.build_query(start_date, end_date, **kwargs)
        query_result = agg_query.execute().to_dict()
        res = self.process_query_result(query_result, start_date, end_date)
        return res
<|endoftext|>