<commit_before>from django.conf import settings
import hashlib
import json
from pathlib import Path
from .tikalib import tika_parse, extract_meta, tika_lang
from . import emails
from . import text
from . import queues
from . import models
from . import archives
from . import pst
from .content_types import guess_content_type, guess_filetype
from .utils import chunks

def _path_bits(doc):
    if doc.container:
        yield from _path_bits(doc.container)
        if emails.is_email(doc.container):
            yield doc.filename
        else:
            yield doc.path
    else:
        yield doc.path

def _calculate_hashes(opened_file):

    md5 = hashlib.md5()
    sha1 = hashlib.sha1()

    for data in chunks(opened_file):
        md5.update(data)
        sha1.update(data)

    fsize = opened_file.tell()

    return (md5.hexdigest(), sha1.hexdigest(), fsize)

def digest(doc):
    if not doc.sha1:
        with doc.open() as f:
            md5, sha1, fsize = _calculate_hashes(f)
        if not doc.disk_size:
            doc.disk_size = fsize
        doc.sha1 = sha1
        doc.md5 = md5
        doc.save()

    data = {
        'path': '//'.join(_path_bits(doc)),
        'lang': None,
        'sha1': doc.sha1,
        'md5': doc.md5,
        'filename': doc.filename,
        'rev': doc.rev,
    }

    if emails.is_email(doc):
        data.update(emails.parse_email(doc))

    if doc.container_id:
        data['message'] = doc.container_id

    filetype = guess_filetype(doc)
    data['type'] = filetype

    if filetype in ['text', 'html']:
        data['text'] = text.get_text(doc)

    if filetype in settings.TIKA_FILE_TYPES and doc.disk_size <= settings.MAX_TIKA_FILE_SIZE:
        parsed = tika_parse(doc.sha1, doc.open)
        data['text'] = (parsed.get('content') or '').strip()
        data.update(extract_meta(parsed.get('metadata', {})))

    if settings.SNOOP_ANALYZE_LANG:
        if 'text' in data and len(data['text']) > 100:
            data['lang'] = tika_lang(data['text'])[:2]

    ocr_items = list(models.Ocr.objects.filter(md5=doc.md5))
    if ocr_items:
        data['ocr'] = {ocr.tag: ocr.text for ocr in ocr_items}

    if archives.is_archive(doc):
        data.update(archives.list_files(doc))
    elif pst.is_pst_file(doc):
        data.update(pst.list_files(doc))

    return data

def create_children(doc, data, verbose=True):
    children_info = []
    if emails.is_email(doc):
        for name, info in data.get('attachments', {}).items():
            children_info.append({
                'path': name,
                'content_type': info['content_type'],
                'filename': info['filename'],
                'size': info['size'],
            })
    elif archives.is_archive(doc) or pst.is_pst_file(doc):
        for path in data['file_list']:
            children_info.append({
                'path': path,
                'content_type': guess_content_type(path),
                'filename': Path(path).name,
                'size': 0,
            })
        for path in data['folder_list']:
            children_info.append({
                'path': path,
                'content_type': 'application/x-directory',
                'filename': Path(path).name,
                'size': 0,
            })

    for info in children_info:
        child, created = models.Document.objects.update_or_create(
            container=doc,
            path=info['path'],
            defaults={
                'disk_size': info['size'],
                'content_type': info['content_type'],
                'filename': info['filename'],
            },
        )

        if created:
            queues.put('digest', {'id': child.id}, verbose=verbose)
            if verbose: print('new child', child.id)

def worker(id, verbose):
    try:
        document = models.Document.objects.get(id=id)
    except models.Document.DoesNotExist:
        if verbose: print('MISSING')
        return

    try:
        data = digest(document)

    except emails.MissingEmlxPart:
        document.broken = 'missing_emlx_part'
        document.save()
        if verbose: print('missing_emlx_part')
        return

    except emails.PayloadError:
        document.broken = 'payload_error'
        document.save()
        if verbose: print('payload_error')
        return

    except emails.CorruptedFile:
        document.broken = 'corrupted_file'
        document.save()
        if verbose: print('corrupted_file')
        return

    except archives.EncryptedArchiveFile:
        document.broken = 'encrypted archive'
        document.save()
        if verbose: print('encrypted archive')
        return

    except archives.MissingArchiveFile:
        document.broken = 'missing archive file'
        document.save()
        if verbose: print('missing archive file')
        return

    except archives.ExtractingFailed:
        document.broken = 'extracting archive with 7z failed'
        document.save()
        if verbose: print('extracting archive with 7z failed')
        return

    except pst.MissingPSTFile:
        document.broken = 'missing pst file'
        document.save()
        if verbose: print('missing pst file')
        return

    except pst.PSTExtractionFailed:
        document.broken = 'extracting archive with readpst failed'
        document.save()
        if verbose: print('extracting archive with readpst failed')
        return

    else:
        if document.broken:
            if verbose: print('removing broken flag', document.broken)
            document.broken = ''
            document.save()

    create_children(document, data, verbose)

    models.Digest.objects.update_or_create(
        id=document.id,
        defaults={'data': json.dumps(data)},
    )

    if verbose: print('type:', data.get('type'))

    queues.put('index', {'id': document.id}, verbose=verbose)
<commit_msg>mark gpg failure as a broken document<commit_after>from django.conf import settings
import hashlib
import json
from pathlib import Path
from .tikalib import tika_parse, extract_meta, tika_lang
from . import emails
from . import text
from . import queues
from . import models
from . import archives
from . import pst
from . import pgp
from .content_types import guess_content_type, guess_filetype
from .utils import chunks

def _path_bits(doc):
    if doc.container:
        yield from _path_bits(doc.container)
        if emails.is_email(doc.container):
            yield doc.filename
        else:
            yield doc.path
    else:
        yield doc.path

def _calculate_hashes(opened_file):

    md5 = hashlib.md5()
    sha1 = hashlib.sha1()

    for data in chunks(opened_file):
        md5.update(data)
        sha1.update(data)

    fsize = opened_file.tell()

    return (md5.hexdigest(), sha1.hexdigest(), fsize)

def digest(doc):
    if not doc.sha1:
        with doc.open() as f:
            md5, sha1, fsize = _calculate_hashes(f)
        if not doc.disk_size:
            doc.disk_size = fsize
        doc.sha1 = sha1
        doc.md5 = md5
        doc.save()

    data = {
        'path': '//'.join(_path_bits(doc)),
        'lang': None,
        'sha1': doc.sha1,
        'md5': doc.md5,
        'filename': doc.filename,
        'rev': doc.rev,
    }

    if emails.is_email(doc):
        data.update(emails.parse_email(doc))

    if doc.container_id:
        data['message'] = doc.container_id

    filetype = guess_filetype(doc)
    data['type'] = filetype

    if filetype in ['text', 'html']:
        data['text'] = text.get_text(doc)

    if filetype in settings.TIKA_FILE_TYPES and doc.disk_size <= settings.MAX_TIKA_FILE_SIZE:
        parsed = tika_parse(doc.sha1, doc.open)
        data['text'] = (parsed.get('content') or '').strip()
        data.update(extract_meta(parsed.get('metadata', {})))

    if settings.SNOOP_ANALYZE_LANG:
        if 'text' in data and len(data['text']) > 100:
            data['lang'] = tika_lang(data['text'])[:2]

    ocr_items = list(models.Ocr.objects.filter(md5=doc.md5))
    if ocr_items:
        data['ocr'] = {ocr.tag: ocr.text for ocr in ocr_items}

    if archives.is_archive(doc):
        data.update(archives.list_files(doc))
    elif pst.is_pst_file(doc):
        data.update(pst.list_files(doc))

    return data

def create_children(doc, data, verbose=True):
    children_info = []
    if emails.is_email(doc):
        for name, info in data.get('attachments', {}).items():
            children_info.append({
                'path': name,
                'content_type': info['content_type'],
                'filename': info['filename'],
                'size': info['size'],
            })
    elif archives.is_archive(doc) or pst.is_pst_file(doc):
        for path in data['file_list']:
            children_info.append({
                'path': path,
                'content_type': guess_content_type(path),
                'filename': Path(path).name,
                'size': 0,
            })
        for path in data['folder_list']:
            children_info.append({
                'path': path,
                'content_type': 'application/x-directory',
                'filename': Path(path).name,
                'size': 0,
            })

    for info in children_info:
        child, created = models.Document.objects.update_or_create(
            container=doc,
            path=info['path'],
            defaults={
                'disk_size': info['size'],
                'content_type': info['content_type'],
                'filename': info['filename'],
            },
        )

        if created:
            queues.put('digest', {'id': child.id}, verbose=verbose)
            if verbose: print('new child', child.id)

def worker(id, verbose):
    try:
        document = models.Document.objects.get(id=id)
    except models.Document.DoesNotExist:
        if verbose: print('MISSING')
        return

    try:
        data = digest(document)

    except emails.MissingEmlxPart:
        document.broken = 'missing_emlx_part'
        document.save()
        if verbose: print('missing_emlx_part')
        return

    except emails.PayloadError:
        document.broken = 'payload_error'
        document.save()
        if verbose: print('payload_error')
        return

    except emails.CorruptedFile:
        document.broken = 'corrupted_file'
        document.save()
        if verbose: print('corrupted_file')
        return

    except archives.EncryptedArchiveFile:
        document.broken = 'encrypted archive'
        document.save()
        if verbose: print('encrypted archive')
        return

    except archives.MissingArchiveFile:
        document.broken = 'missing archive file'
        document.save()
        if verbose: print('missing archive file')
        return

    except archives.ExtractingFailed:
        document.broken = 'extracting archive with 7z failed'
        document.save()
        if verbose: print('extracting archive with 7z failed')
        return

    except pst.MissingPSTFile:
        document.broken = 'missing pst file'
        document.save()
        if verbose: print('missing pst file')
        return

    except pst.PSTExtractionFailed:
        document.broken = 'extracting archive with readpst failed'
        document.save()
        if verbose: print('extracting archive with readpst failed')
        return

    except pgp.DecryptionError:
        document.broken = 'decryption failed'
        document.save()
        if verbose: print('decryption failed')
        return

    else:
        if document.broken:
            if verbose: print('removing broken flag', document.broken)
            document.broken = ''
            document.save()

    create_children(document, data, verbose)

    models.Digest.objects.update_or_create(
        id=document.id,
        defaults={'data': json.dumps(data)},
    )

    if verbose: print('type:', data.get('type'))

    queues.put('index', {'id': document.id}, verbose=verbose)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Prepares the WikiProject category tree
Copyright (C) 2015 James Hare
Licensed under MIT License: http://mitlicense.org
"""

from pprint import pprint
from project_index import WikiProjectTools

class WikiProjectCategories:
    def generate(self, audit=False, production=True):
        '''
        "audit" prints the category tree to stdout
        "production" returns a dictionary for use in some other script
        '''

        wptools = WikiProjectTools()

        tree = {}
        category = 'WikiProjects_by_area' # seed category

        query = wptools.query('wiki', 'select distinct page.page_title from categorylinks join page on categorylinks.cl_from=page.page_id where page_namespace = 14 and cl_to = "WikiProjects_by_area" and page_title like "%\_WikiProjects";', None)
        for row in query:
            while len(query) > 0:
                category = row[0].decode('utf-8')
                tree[category] = {}
                query = 'select distinct page.page_title from categorylinks join page on categorylinks.cl_from=page.page_id where page_namespace = 14 and cl_to = "{0}" and page_title like "%\_WikiProjects";'.format(category)

        if audit == True:
            pprint(tree)

        if production == True:
            return tree

if __name__ == "__main__":
    wpc = WikiProjectCategories()
    wpc.generate(audit=True, production=False)  # If the script is invoked directly, run in audit mode<commit_msg>Thank you Merljn!<commit_after># -*- coding: utf-8 -*-
"""
Prepares the WikiProject category tree
Copyright (C) 2015 James Hare, Merlijn Van Deen
Licensed under MIT License: http://mitlicense.org
"""

from pprint import pprint
from project_index import WikiProjectTools

def build_cat_tree(cat_name, max_depth=5):
    if max_depth == 0:
        return None
    query = wptools.query('wiki', 'select distinct page.page_title from categorylinks join page on categorylinks.cl_from=page.page_id where page_namespace = 14 and cl_to = "{0}" and page_title like "%\_WikiProjects";', cat_name)
    retval = {}
    for row in query:
        category = row[0].decode('utf-8')
        retval[category] = build_cat_tree(category, max_depth=max_depth-1)
    return retval

class WikiProjectCategories:
    def generate(self, audit=False, production=True):
        '''
        "audit" prints the category tree to stdout
        "production" returns a dictionary for use in some other script
        '''

        wptools = WikiProjectTools()
        tree = build_cat_tree('WikiProjects_by_area', max_depth=2)

        if audit == True:
            pprint(tree)

        if production == True:
            return tree

if __name__ == "__main__":
    wpc = WikiProjectCategories()
    wpc.generate(audit=True, production=False)  # If the script is invoked directly, run in audit mode

<|endoftext|>


<commit_before>from django.utils import timezone
from django.db.models import Max

from breach.analyzer import decide_next_world_state
from breach.models import SampleSet, Round
from breach.sniffer import Sniffer

import string
import requests
import logging

logger = logging.getLogger(__name__)

SAMPLES_PER_SAMPLESET = 100

class Strategy(object):
    def __init__(self, victim):
        self._victim = victim
        self._sniffer = Sniffer(victim.snifferendpoint)

        # Extract maximum round index for the current victim.
        current_round_index = Round.objects.filter(victim=self._victim).aggregate(Max('index'))['index__max']

        if not current_round_index:
            current_round_index = 1
            self._analyzed = True
            self.begin_attack()

        self._round = Round.objects.filter(victim=self._victim, index=current_round_index)[0]
        self._analyzed = False

    def _build_candidates_divide_conquer(self, state):
        candidate_alphabet_cardinality = len(state['knownalphabet']) / 2

        bottom_half = state['knownalphabet'][:candidate_alphabet_cardinality]
        top_half = state['knownalphabet'][candidate_alphabet_cardinality:]

        return [bottom_half, top_half]

    def _build_candidates_serial(self, state):
        return state['knownalphabet']

    def _build_candidates(self, state):
        '''Given a state of the world, produce a list of candidate alphabets.'''
        methods = {
            'serial': self._build_candidates_serial,
            'divide&conquer': self._build_candidates_divide_conquer
        }
        return methods[self._victim.method](state)

    def _get_first_round_state(self):
        return {
            'knownsecret': self._victim.target.prefix,
            'candidatealphabet': self._victim.target.alphabet,
            'knownalphabet': self._victim.target.alphabet
        }

    def _get_unstarted_samplesets(self):
        return SampleSet.objects.filter(
            round=self._round,
            started=None
        )

    def _reflection(self, sampleset):
        # We use '^' as a separator symbol and we assume it is not part of the
        # secret. We also assume it will not be in the content.

        sentinel = '^'

        assert(sentinel not in self._round.knownalphabet)
        alphabet_complement = list(set(string.punctuation + string.ascii_letters + string.digits) - set(self._round.knownalphabet))

        huffman_complement = set(self._round.knownalphabet) - set(sampleset.candidatealphabet)

        candidate_secrets = set()

        for letter in sampleset.candidatealphabet:
            candidate_secret = self._round.knownsecret + letter
            candidate_secrets.add(candidate_secret)

        # Add as many dummy symbols as necessary, so that all different candidate
        # alphabets have the same amount of total data.
        candidate_balance = self._round.roundcardinality - len(candidate_secrets)
        assert(len(alphabet_complement) > candidate_balance)
        candidate_balance = [self._round.knownsecret + c for c in alphabet_complement[0:candidate_balance]]

        huffman_balance = self._round.roundcardinality - len(huffman_complement)
        assert(len(alphabet_complement) > huffman_balance)
        huffman_balance = alphabet_complement[0:huffman_balance]

        reflected_data = [
            '',
            sentinel.join(list(candidate_secrets) + candidate_balance),
            sentinel.join(list(huffman_complement) + huffman_balance),
            ''
        ]

        reflection = sentinel.join(reflected_data)

        return reflection

    def _sampleset_to_work(self, sampleset):
        return {
            'url': self._victim.target.endpoint % self._reflection(sampleset),
            'amount': SAMPLES_PER_SAMPLESET,
            'timeout': 0
        }

    def get_work(self):
        '''Produces work for the victim.

        Pre-condition: There is already work to do.'''

        try:
            self._sniffer.start(self._victim.sourceip, self._victim.target.host)
        except (requests.HTTPError, requests.exceptions.ConnectionError), err:
            if isinstance(err, requests.HTTPError):
                status_code = err.response.status_code
                logger.warning('Caught {} while trying to start sniffer.'.format(status_code))

                # If status was raised due to conflict,
                # delete already existing sniffer.
                if status_code == 409:
                    try:
                        self._sniffer.delete(self._victim.sourceip, self._victim.target.host)
                    except (requests.HTTPError, requests.exceptions.ConnectionError), err:
                        logger.warning('Caught error when trying to delete sniffer: {}'.format(err))

            elif isinstance(err, requests.exceptions.ConnectionError):
                logger.warning('Caught ConnectionError')

            return {}

        unstarted_samplesets = self._get_unstarted_samplesets()

        assert(unstarted_samplesets)

        sampleset = unstarted_samplesets[0]
        sampleset.started = timezone.now()
        sampleset.save()

        work = self._sampleset_to_work(sampleset)

        logger.debug('Giving work: {}'.format(work))

        return work

    def _get_current_sampleset(self):
        started_samplesets = SampleSet.objects.filter(round=self._round, completed=None).exclude(started=None)

        assert(len(started_samplesets) == 1)

        sampleset = started_samplesets[0]

        return sampleset

    def _mark_current_work_completed(self, capture=None):
        sampleset = self._get_current_sampleset()
        sampleset.completed = timezone.now()

        if capture:
            sampleset.success = True
            sampleset.data = capture
        else:
            # Sampleset data collection failed,
            # create a new sampleset for the same attack element
            s = SampleSet(
                round=sampleset.round,
                candidatealphabet=sampleset.candidatealphabet
            )
            s.save()

        sampleset.save()

    def _collect_capture(self):
        captured_data = self._sniffer.read(self._victim.sourceip, self._victim.target.host)
        return captured_data['capture']

    def _analyze_current_round(self):
        '''Analyzes the current round samplesets to extract a decision.'''

        current_round_samplesets = SampleSet.objects.filter(round=self._round, success=True)
        self._decision = decide_next_world_state(current_round_samplesets)

        logger.debug('Decision: {}'.format(self._decision))
        self._analyzed = True

    def _round_is_completed(self):
        '''Checks if current round is completed.'''

        assert(self._analyzed)

        # Do we need to collect more samplesets to build up confidence?
        return self._decision['confidence'] > 1

    def _create_next_round(self):
        assert(self._round_is_completed())

        self._create_round(self._decision['state'])

    def _create_round(self, state):
        '''Creates a new round based on the analysis of the current round.'''

        assert(self._analyzed)

        candidate_alphabets = self._build_candidates(state)

        # This next round could potentially be the final round.
        # A final round has the complete secret stored in knownsecret.
        next_round = Round(
            victim=self._victim,
            index=self._round.index + 1 if hasattr(self, '_round') else 1,
            roundcardinality=max(map(len, candidate_alphabets)),
            amount=SAMPLES_PER_SAMPLESET,
            knownalphabet=state['knownalphabet'],
            knownsecret=state['knownsecret']
        )
        next_round.save()

        self._round = next_round

    def _create_round_samplesets(self):
        state = {
            'knownalphabet': self._round.knownalphabet,
            'knownsecret': self._round.knownsecret
        }

        candidate_alphabets = self._build_candidates(state)

        for candidate in candidate_alphabets:
            sampleset = SampleSet(
                round=self._round,
                candidatealphabet=candidate
            )
            sampleset.save()

    def _attack_is_completed(self):
        return len(self._round.knownsecret) == self._victim.target.secretlength

    def work_completed(self, success=True):
        '''Receives and consumes work completed from the victim, analyzes
        the work, and returns True if the attack is complete (victory),
        otherwise returns False if more work is needed.

        It also creates the new work that is needed.

        Post-condition: Either the attack is completed, or there is work to
        do (there are unstarted samplesets in the database).'''

        try:
            # Call sniffer to get captured data
            capture = self._collect_capture()
            logger.debug('Collected capture with length: {}'.format(len(capture)))

            # Stop data collection and delete sniffer
            self._sniffer.delete(self._victim.sourceip, self._victim.target.host)
        except (requests.HTTPError, requests.exceptions.ConnectionError), err:
            if isinstance(err, requests.HTTPError):
                status_code = err.response.status_code
                logger.warning('Caught {} while trying to collect capture and delete sniffer.'.format(status_code))

                # If status was raised due to malformed capture,
                # delete sniffer to avoid conflict.
                if status_code == 422:
                    try:
                        self._sniffer.delete(self._victim.sourceip, self._victim.target.host)
                    except (requests.HTTPError, requests.exceptions.ConnectionError), err:
                        logger.warning('Caught error when trying to delete sniffer: {}'.format(err))

            elif isinstance(err, requests.exceptions.ConnectionError):
                logger.warning('Caught ConnectionError')

            self._mark_current_work_completed()
            return False

        self._mark_current_work_completed(capture)

        round_samplesets = SampleSet.objects.filter(round=self._round)
        unstarted_samplesets = round_samplesets.filter(started=None)

        if unstarted_samplesets:
            # Batch is not yet complete, we need to collect more samplesets
            # that have already been created for this batch.
            return False

        # All batches are completed.
        self._analyze_current_round()

        if self._round_is_completed():
            # Advance to the next round.
            self._create_next_round()

            if self._attack_is_completed():
                return True

        # Not enough confidence, we need to create more samplesets to be
        # collected for this round.
        self._create_round_samplesets()

        return False

    def begin_attack(self):
        self._create_round(self._get_first_round_state())
        self._create_round_samplesets()
<commit_msg>Abort sampleset if realtime posted unsuccessful<commit_after>from django.utils import timezone
from django.db.models import Max

from breach.analyzer import decide_next_world_state
from breach.models import SampleSet, Round
from breach.sniffer import Sniffer

import string
import requests
import logging

logger = logging.getLogger(__name__)

SAMPLES_PER_SAMPLESET = 100

class Strategy(object):
    def __init__(self, victim):
        self._victim = victim
        self._sniffer = Sniffer(victim.snifferendpoint)

        # Extract maximum round index for the current victim.
        current_round_index = Round.objects.filter(victim=self._victim).aggregate(Max('index'))['index__max']

        if not current_round_index:
            current_round_index = 1
            self._analyzed = True
            self.begin_attack()

        self._round = Round.objects.filter(victim=self._victim, index=current_round_index)[0]
        self._analyzed = False

    def _build_candidates_divide_conquer(self, state):
        candidate_alphabet_cardinality = len(state['knownalphabet']) / 2

        bottom_half = state['knownalphabet'][:candidate_alphabet_cardinality]
        top_half = state['knownalphabet'][candidate_alphabet_cardinality:]

        return [bottom_half, top_half]

    def _build_candidates_serial(self, state):
        return state['knownalphabet']

    def _build_candidates(self, state):
        '''Given a state of the world, produce a list of candidate alphabets.'''
        methods = {
            'serial': self._build_candidates_serial,
            'divide&conquer': self._build_candidates_divide_conquer
        }
        return methods[self._victim.method](state)

    def _get_first_round_state(self):
        return {
            'knownsecret': self._victim.target.prefix,
            'candidatealphabet': self._victim.target.alphabet,
            'knownalphabet': self._victim.target.alphabet
        }

    def _get_unstarted_samplesets(self):
        return SampleSet.objects.filter(
            round=self._round,
            started=None
        )

    def _reflection(self, sampleset):
        # We use '^' as a separator symbol and we assume it is not part of the
        # secret. We also assume it will not be in the content.

        sentinel = '^'

        assert(sentinel not in self._round.knownalphabet)
        alphabet_complement = list(set(string.punctuation + string.ascii_letters + string.digits) - set(self._round.knownalphabet))

        huffman_complement = set(self._round.knownalphabet) - set(sampleset.candidatealphabet)

        candidate_secrets = set()

        for letter in sampleset.candidatealphabet:
            candidate_secret = self._round.knownsecret + letter
            candidate_secrets.add(candidate_secret)

        # Add as many dummy symbols as necessary, so that all different candidate
        # alphabets have the same amount of total data.
        candidate_balance = self._round.roundcardinality - len(candidate_secrets)
        assert(len(alphabet_complement) > candidate_balance)
        candidate_balance = [self._round.knownsecret + c for c in alphabet_complement[0:candidate_balance]]

        huffman_balance = self._round.roundcardinality - len(huffman_complement)
        assert(len(alphabet_complement) > huffman_balance)
        huffman_balance = alphabet_complement[0:huffman_balance]

        reflected_data = [
            '',
            sentinel.join(list(candidate_secrets) + candidate_balance),
            sentinel.join(list(huffman_complement) + huffman_balance),
            ''
        ]

        reflection = sentinel.join(reflected_data)

        return reflection

    def _sampleset_to_work(self, sampleset):
        return {
            'url': self._victim.target.endpoint % self._reflection(sampleset),
            'amount': SAMPLES_PER_SAMPLESET,
            'timeout': 0
        }

    def get_work(self):
        '''Produces work for the victim.

        Pre-condition: There is already work to do.'''

        try:
            self._sniffer.start(self._victim.sourceip, self._victim.target.host)
        except (requests.HTTPError, requests.exceptions.ConnectionError), err:
            if isinstance(err, requests.HTTPError):
                status_code = err.response.status_code
                logger.warning('Caught {} while trying to start sniffer.'.format(status_code))

                # If status was raised due to conflict,
                # delete already existing sniffer.
                if status_code == 409:
                    try:
                        self._sniffer.delete(self._victim.sourceip, self._victim.target.host)
                    except (requests.HTTPError, requests.exceptions.ConnectionError), err:
                        logger.warning('Caught error when trying to delete sniffer: {}'.format(err))

            elif isinstance(err, requests.exceptions.ConnectionError):
                logger.warning('Caught ConnectionError')

            return {}

        unstarted_samplesets = self._get_unstarted_samplesets()

        assert(unstarted_samplesets)

        sampleset = unstarted_samplesets[0]
        sampleset.started = timezone.now()
        sampleset.save()

        work = self._sampleset_to_work(sampleset)

        logger.debug('Giving work: {}'.format(work))

        return work

    def _get_current_sampleset(self):
        started_samplesets = SampleSet.objects.filter(round=self._round, completed=None).exclude(started=None)

        assert(len(started_samplesets) == 1)

        sampleset = started_samplesets[0]

        return sampleset

    def _mark_current_work_completed(self, capture=None):
        sampleset = self._get_current_sampleset()
        sampleset.completed = timezone.now()

        if capture:
            sampleset.success = True
            sampleset.data = capture
        else:
            # Sampleset data collection failed,
            # create a new sampleset for the same attack element
            s = SampleSet(
                round=sampleset.round,
                candidatealphabet=sampleset.candidatealphabet
            )
            s.save()

        sampleset.save()

    def _collect_capture(self):
        captured_data = self._sniffer.read(self._victim.sourceip, self._victim.target.host)
        return captured_data['capture']

    def _analyze_current_round(self):
        '''Analyzes the current round samplesets to extract a decision.'''

        current_round_samplesets = SampleSet.objects.filter(round=self._round, success=True)
        self._decision = decide_next_world_state(current_round_samplesets)

        logger.debug('Decision: {}'.format(self._decision))
        self._analyzed = True

    def _round_is_completed(self):
        '''Checks if current round is completed.'''

        assert(self._analyzed)

        # Do we need to collect more samplesets to build up confidence?
        return self._decision['confidence'] > 1

    def _create_next_round(self):
        assert(self._round_is_completed())

        self._create_round(self._decision['state'])

    def _create_round(self, state):
        '''Creates a new round based on the analysis of the current round.'''

        assert(self._analyzed)

        candidate_alphabets = self._build_candidates(state)

        # This next round could potentially be the final round.
        # A final round has the complete secret stored in knownsecret.
        next_round = Round(
            victim=self._victim,
            index=self._round.index + 1 if hasattr(self, '_round') else 1,
            roundcardinality=max(map(len, candidate_alphabets)),
            amount=SAMPLES_PER_SAMPLESET,
            knownalphabet=state['knownalphabet'],
            knownsecret=state['knownsecret']
        )
        next_round.save()

        self._round = next_round

    def _create_round_samplesets(self):
        state = {
            'knownalphabet': self._round.knownalphabet,
            'knownsecret': self._round.knownsecret
        }

        candidate_alphabets = self._build_candidates(state)

        for candidate in candidate_alphabets:
            sampleset = SampleSet(
                round=self._round,
                candidatealphabet=candidate
            )
            sampleset.save()

    def _attack_is_completed(self):
        return len(self._round.knownsecret) == self._victim.target.secretlength

    def work_completed(self, success=True):
        '''Receives and consumes work completed from the victim, analyzes
        the work, and returns True if the attack is complete (victory),
        otherwise returns False if more work is needed.

        It also creates the new work that is needed.

        Post-condition: Either the attack is completed, or there is work to
        do (there are unstarted samplesets in the database).'''
        try:
            if success:
                # Call sniffer to get captured data
                capture = self._collect_capture()
                logger.debug('Collected capture with length: {}'.format(len(capture)))
            else:
                logger.debug('Client returned fail to realtime')
                assert success

            # Stop data collection and delete sniffer
            self._sniffer.delete(self._victim.sourceip, self._victim.target.host)
        except (requests.HTTPError, requests.exceptions.ConnectionError, AssertionError), err:
            if isinstance(err, requests.HTTPError):
                status_code = err.response.status_code
                logger.warning('Caught {} while trying to collect capture and delete sniffer.'.format(status_code))

                # If status was raised due to malformed capture,
                # delete sniffer to avoid conflict.
                if status_code == 422:
                    try:
                        self._sniffer.delete(self._victim.sourceip, self._victim.target.host)
                    except (requests.HTTPError, requests.exceptions.ConnectionError), err:
                        logger.warning('Caught error when trying to delete sniffer: {}'.format(err))

            elif isinstance(err, requests.exceptions.ConnectionError):
                logger.warning('Caught ConnectionError')

            elif isinstance(err, AssertionError):
                logger.warning('Realtime reported unsuccessful capture')
                try:
                    self._sniffer.delete(self._victim.sourceip, self._victim.target.host)
                except (requests.HTTPError, requests.exceptions.ConnectionError), err:
                    logger.warning('Caught error when trying to delete sniffer: {}'.format(err))

            self._mark_current_work_completed()
            return False

        self._mark_current_work_completed(capture)

        round_samplesets = SampleSet.objects.filter(round=self._round)
        unstarted_samplesets = round_samplesets.filter(started=None)

        if unstarted_samplesets:
            # Batch is not yet complete, we need to collect more samplesets
            # that have already been created for this batch.
            return False

        # All batches are completed.
        self._analyze_current_round()

        if self._round_is_completed():
            # Advance to the next round.
            self._create_next_round()

            if self._attack_is_completed():
                return True

        # Not enough confidence, we need to create more samplesets to be
        # collected for this round.
        self._create_round_samplesets()

        return False

    def begin_attack(self):
        self._create_round(self._get_first_round_state())
        self._create_round_samplesets()
<|endoftext|>


<commit_before>import zmq
import asyncio
from tornado.ioloop import IOLoop
import capnp
import os
from app.jobs import Job

CWD = os.path.dirname(__file__)
tidings_capnp = capnp.load(os.path.join(CWD, 'capnp/src/tidings.capnp'))
plumbing_capnp = capnp.load(os.path.join(CWD, 'capnp/src/plumbing.capnp'))

class QueueClient:
    def __init__(self, ioloop, host, port):
        self.ctx = zmq.asyncio.Context()
        self.socket = self.ctx.socket(zmq.REQ)
        self.socket.connect('{}:{}'.format(host, port))
        ioloop.spawn_callback(self._get_stats_loop)
        self.submitted = {}

    def _process_submit_answer(self, resp):
        message = plumbing_capnp.Plumbing.from_bytes_packed(resp[0])
        self.submitted[message.id] = message

    async def _send_request(self, message):
        await self.socket.send_multipart(message)
        return await self.socket.recv_multipart()

    async def _get_stats_loop(self):
        pass

    async def submit_job(self, job: Job):
        message = tidings_capnp.Tiding.new_message()  # TODO finish this
        resp = await self._send_request(message)
        return self._process_submit_answer(resp)

    def get_results(id):
        pass

<commit_msg>Kafka client implementation<commit_after>import logging
import msgpack

from aiokafka import AIOKafkaProducer, AIOKafkaConsumer
from app.jobs import Job

log = logging.getLogger(__name__)

class QueueClient:

    @classmethod
    async def create(cls, ioloop, conf):
        self = QueueClient()
        self.conf = conf
        print('Starting client with config: ', conf)
        self.consumer = AIOKafkaConsumer(
            *conf['kafka.consumer_topics'],
            loop=ioloop, bootstrap_servers=conf['kafka.servers'],
            group_id=conf['kafka.consumer_group'],
            key_deserializer=QueueClient._KeyDeserializer(),
            value_deserializer=QueueClient._ValueDeserializer())

        self.producer = AIOKafkaProducer(
            loop=ioloop, bootstrap_servers=conf['kafka.servers'],
            key_serializer=QueueClient._KeySerializer(),
            value_serializer=QueueClient._ValueSerializer())

        await self.consumer.start()
        await self.producer.start()

        return self

    def __init__(self):
        self.submitted = {}

    async def _consume(self):
        try:
            # Consume messages
            async for msg in self.consumer:
                log.debug("consumed: ", msg.topic, msg.partition, msg.offset,
                          msg.key, msg.value, msg.timestamp)
                if(msg.topic == 'csaopt.data.in.t'):
                    pass
                if(msg.topic == 'csaopt.stats.in.t'):
                    pass
                print(msg)
        finally:
            # Will leave consumer group; perform autocommit if enabled.
            await self.consumer.stop()
            await self.producer.stop()

    async def _send_one(self, topic, key=None, value=None):
        try:
            # Produce message
            await self.producer.send(topic, key=key, value=value)
        finally:
            # Wait for all pending messages to be delivered or expire.
            pass

    async def submit_job(self, job: Job):
        self.submitted[job.id] = job

    def get_results(id):
        pass

    class _KeySerializer:
        def __call__(self, key: str) -> bytes:
            return key.encode(encoding='utf_8') if key is not None else b''

    class _ValueSerializer:
        def __call__(self, value: object) -> bytes:
            return msgpack.packb(value, use_bin_type=True)

    class _KeyDeserializer:
        def __call__(self, keyBytes: bytes) -> str:
            return keyBytes.decode(encoding='utf-8') if bytes is not None else ''

    class _ValueDeserializer:
        def __call__(self, valueBytes: bytes) -> object:
            return msgpack.unpackb(valueBytes)
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
Function to load data from CKAN object classes.

'''
from app.classes.user import User
from app.classes.country import Country
from app.classes.dataset import Dataset
from app.classes.revision import Revision
from app.classes.resource import Resource
from app.classes.gallery_item import GalleryItem

def fetchClassData(key=None, id=None):
  '''
  Loads data from specified CKAN object class.

  '''
  classes = {
      'user': User(id),
      'country': Country(id),
      'dataset': Dataset(id),
      'revision': Revision(id),
      'resource': Resource(id),
      'gallery_item': GalleryItem(id)
    }

  return { 'job': None, 'queue': None, 'data': classes.get(key).info() }
<commit_msg>added filter for selecting only fields determined in config file<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
Function to load data from CKAN object classes.

'''
import app.utilities.load as Load

from app.classes.user import User
from app.classes.country import Country
from app.classes.dataset import Dataset
from app.classes.revision import Revision
from app.classes.resource import Resource
from app.classes.gallery_item import GalleryItem

config = Load.loadJSONFile('config/dev.json')

def _fields(config, key):
  '''
  Extract field names from the database
  property in the configuration file.

  '''
  result = None
  for table in config['database']:
    if table['name'] == key:
      result = [ i['field_name'] for i in table['columns'] ]

  return result

def fetchClassData(key=None, id=None):
  '''
  Loads data from specified CKAN object class.

  '''
  classes = {
      'user': User(id),
      'country': Country(id),
      'dataset': Dataset(id),
      'revision': Revision(id),
      'resource': Resource(id),
      'gallery_item': GalleryItem(id)
    }

  #
  # Selects only the fields
  # of interest from the dictionary.
  #
  result = { k: classes[key].info() for k in _fields(config, key) }
  return result
<|endoftext|>


<commit_before>from approxeng.input import CentredAxis, TriggerAxis, Button, Controller, BinaryAxis
from approxeng.input.sys import read_power_level

XB1S_VENDOR_ID = 1118
XB1S_WIRED_PRODUCT_ID = 746
XB1S_WIRELESS_PRODUCT_ID = 765
XB1S_WIRELESS_PRODUCT_ID_2 = 736

class WiredXBoxOneSPad(Controller):
    """
    Wired XBox One controller, tested with the newer controllers that use bluetooth and are supplied with the XBox One S
    but may work with older versions. Note that the codes and axis mappings for the wired controller are not the same
    as for the wireless!
    """

    def __init__(self, dead_zone=0.1, hot_zone=0.05):
        """
        Create a new xbox one s controller instance

        :param float dead_zone:
            Used to set the dead zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        :param float hot_zone:
            Used to set the hot zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        """
        super(WiredXBoxOneSPad, self).__init__(vendor_id=XB1S_VENDOR_ID,
                                               product_id=XB1S_WIRED_PRODUCT_ID,
                                               controls=[
                                                   Button("X", 307, sname='square'),
                                                   Button("Y", 308, sname='triangle'),
                                                   Button("B", 305, sname='circle'),
                                                   Button("A", 304, sname='cross'),
                                                   Button("Right Stick", 318, sname='rs'),
                                                   Button("Left Stick", 317, sname='ls'),
                                                   Button("View", 314, sname='select'),
                                                   Button("Menu", 315, sname='start'),
                                                   Button("XBox", 316, sname='home'),
                                                   Button("LB", 310, sname='l1'),
                                                   Button("RB", 311, sname='r1'),
                                                   CentredAxis("Left Horizontal", -32768, 32768, 0, sname='lx'),
                                                   CentredAxis("Left Vertical", -32768, 32768, 1, invert=True,
                                                               sname='ly'),
                                                   CentredAxis("Right Horizontal", -32768, 32768, 3, sname='rx'),
                                                   CentredAxis("Right Vertical", -32768, 32768, 4, invert=True,
                                                               sname='ry'),
                                                   TriggerAxis("Left Trigger", 0, 1023, 2, sname='lt',
                                                               button_sname='l2', button_trigger_value=0.2),
                                                   TriggerAxis("Right Trigger", 0, 1023, 5, sname='rt',
                                                               button_sname='r2', button_trigger_value=0.2),
                                                   BinaryAxis("D-pad Horizontal", 16, b1name='dleft', b2name='dright'),
                                                   BinaryAxis("D-pad Vertical", 17, b1name='dup', b2name='ddown')
                                               ],
                                               dead_zone=dead_zone,
                                               hot_zone=hot_zone)

    def __repr__(self):
        return 'Wired Microsoft XBox One S controller'

class WirelessXBoxOneSPad(Controller):
    """
    Wireless over bluetooth - for some reason the control ranges on this are different to when wired and all the buttons
    are different.

    To make the wireless pair work with linux you need to set a kernel module option on the bluetooth module. This can
    be done by creating a file in /etc/modprobe.d with the line 'options bluetooth disable_ertm=Y' and rebooting.
    """

    def __init__(self, dead_zone=0.1, hot_zone=0.05):
        """
        Create a new xbox one s controller instance

        :param float dead_zone:
            Used to set the dead zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        :param float hot_zone:
            Used to set the hot zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        """
        super(WirelessXBoxOneSPad, self).__init__(vendor_id=XB1S_VENDOR_ID,
                                                  product_id=XB1S_WIRELESS_PRODUCT_ID,
                                                  controls=[
                                                      Button("X", 307, sname='square'),
                                                      Button("Y", 308, sname='triangle'),
                                                      Button("B", 305, sname='circle'),
                                                      Button("A", 304, sname='cross'),
                                                      Button("Right Stick", 318, sname='rs'),
                                                      Button("Left Stick", 317, sname='ls'),
                                                      Button("View", 158, sname='select'),
                                                      Button("Menu", 315, sname='start'),
                                                      Button("XBox", 172, sname='home'),
                                                      Button("LB", 310, sname='l1'),
                                                      Button("RB", 311, sname='r1'),
                                                      CentredAxis("Left Horizontal", 0, 65335, 0, sname='lx'),
                                                      CentredAxis("Left Vertical", 0, 65335, 1, invert=True,
                                                                  sname='ly'),
                                                      CentredAxis("Right Horizontal", 0, 65335, 2, sname='rx'),
                                                      CentredAxis("Right Vertical", 0, 65335, 5, invert=True,
                                                                  sname='ry'),
                                                      TriggerAxis("Left Trigger", 0, 1023, 10, sname='lt',
                                                                  button_sname='l2', button_trigger_value=0.1),
                                                      TriggerAxis("Right Trigger", 0, 1023, 9, sname='rt',
                                                                  button_sname='r2', button_trigger_value=0.1),
                                                      BinaryAxis("D-pad Horizontal", 16, b1name='dleft',
                                                                 b2name='dright'),
                                                      BinaryAxis("D-pad Vertical", 17, b1name='dup', b2name='ddown')
                                                  ],
                                                  dead_zone=dead_zone,
                                                  hot_zone=hot_zone)

    @property
    def battery_level(self):
        return float(read_power_level(self.device_unique_name)) / 100.0

    def __repr__(self):
        return 'Wireless Microsoft XBox One S controller'
<commit_msg>XB1 controller seems to have changed its product key, again...<commit_after>from approxeng.input import CentredAxis, TriggerAxis, Button, Controller, BinaryAxis
from approxeng.input.sys import read_power_level

XB1S_VENDOR_ID = 1118
XB1S_WIRED_PRODUCT_ID = 746
XB1S_WIRELESS_PRODUCT_ID = 765
XB1S_WIRELESS_PRODUCT_ID_2 = 736

class WiredXBoxOneSPad(Controller):
    """
    Wired XBox One controller, tested with the newer controllers that use bluetooth and are supplied with the XBox One S
    but may work with older versions. Note that the codes and axis mappings for the wired controller are not the same
    as for the wireless!
    """

    def __init__(self, dead_zone=0.1, hot_zone=0.05):
        """
        Create a new xbox one s controller instance

        :param float dead_zone:
            Used to set the dead zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        :param float hot_zone:
            Used to set the hot zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        """
        super(WiredXBoxOneSPad, self).__init__(vendor_id=XB1S_VENDOR_ID,
                                               product_id=XB1S_WIRED_PRODUCT_ID,
                                               controls=[
                                                   Button("X", 307, sname='square'),
                                                   Button("Y", 308, sname='triangle'),
                                                   Button("B", 305, sname='circle'),
                                                   Button("A", 304, sname='cross'),
                                                   Button("Right Stick", 318, sname='rs'),
                                                   Button("Left Stick", 317, sname='ls'),
                                                   Button("View", 314, sname='select'),
                                                   Button("Menu", 315, sname='start'),
                                                   Button("XBox", 316, sname='home'),
                                                   Button("LB", 310, sname='l1'),
                                                   Button("RB", 311, sname='r1'),
                                                   CentredAxis("Left Horizontal", -32768, 32768, 0, sname='lx'),
                                                   CentredAxis("Left Vertical", -32768, 32768, 1, invert=True,
                                                               sname='ly'),
                                                   CentredAxis("Right Horizontal", -32768, 32768, 3, sname='rx'),
                                                   CentredAxis("Right Vertical", -32768, 32768, 4, invert=True,
                                                               sname='ry'),
                                                   TriggerAxis("Left Trigger", 0, 1023, 2, sname='lt',
                                                               button_sname='l2', button_trigger_value=0.2),
                                                   TriggerAxis("Right Trigger", 0, 1023, 5, sname='rt',
                                                               button_sname='r2', button_trigger_value=0.2),
                                                   BinaryAxis("D-pad Horizontal", 16, b1name='dleft', b2name='dright'),
                                                   BinaryAxis("D-pad Vertical", 17, b1name='dup', b2name='ddown')
                                               ],
                                               dead_zone=dead_zone,
                                               hot_zone=hot_zone)

    def __repr__(self):
        return 'Wired Microsoft XBox One S controller'

class WirelessXBoxOneSPad(Controller):
    """
    Wireless over bluetooth - for some reason the control ranges on this are different to when wired and all the buttons
    are different.

    To make the wireless pair work with linux you need to set a kernel module option on the bluetooth module. This can
    be done by creating a file in /etc/modprobe.d with the line 'options bluetooth disable_ertm=Y' and rebooting.
    """

    def __init__(self, dead_zone=0.1, hot_zone=0.05):
        """
        Create a new xbox one s controller instance

        :param float dead_zone:
            Used to set the dead zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        :param float hot_zone:
            Used to set the hot zone for each :class:`approxeng.input.CentredAxis` and
            :class:`approxeng.input.TriggerAxis` in the controller.
        """
        super(WirelessXBoxOneSPad, self).__init__(vendor_id=XB1S_VENDOR_ID,
                                                  product_id=XB1S_WIRELESS_PRODUCT_ID,
                                                  controls=[
                                                      Button("X", 306, sname='square'),
                                                      Button("Y", 307, sname='triangle'),
                                                      Button("B", 305, sname='circle'),
                                                      Button("A", 304, sname='cross'),
                                                      Button("Right Stick", 313, sname='rs'),
                                                      Button("Left Stick", 312, sname='ls'),
                                                      Button("View", 310, sname='select'),
                                                      Button("Menu", 311, sname='start'),
                                                      Button("XBox", 139, sname='home'),
                                                      Button("LB", 308, sname='l1'),
                                                      Button("RB", 309, sname='r1'),
                                                      CentredAxis("Left Horizontal", 0, 65335, 0, sname='lx'),
                                                      CentredAxis("Left Vertical", 0, 65335, 1, invert=True,
                                                                  sname='ly'),
                                                      CentredAxis("Right Horizontal", 0, 65335, 3, sname='rx'),
                                                      CentredAxis("Right Vertical", 0, 65335, 4, invert=True,
                                                                  sname='ry'),
                                                      TriggerAxis("Left Trigger", 0, 1023, 2, sname='lt',
                                                                  button_sname='l2', button_trigger_value=0.1),
                                                      TriggerAxis("Right Trigger", 0, 1023, 5, sname='rt',
                                                                  button_sname='r2', button_trigger_value=0.1),
                                                      BinaryAxis("D-pad Horizontal", 16, b1name='dleft',
                                                                 b2name='dright'),
                                                      BinaryAxis("D-pad Vertical", 17, b1name='dup', b2name='ddown')
                                                  ],
                                                  dead_zone=dead_zone,
                                                  hot_zone=hot_zone)

    @property
    def battery_level(self):
        return float(read_power_level(self.device_unique_name)) / 100.0

    def __repr__(self):
        return 'Wireless Microsoft XBox One S controller'
<|endoftext|>


<commit_before># Copyright (C) 2010 Peter Teichman

import logging
import math
import random
import re
import sqlite3

import megahal
import tokenizer

log = logging.getLogger("hal")

# use an empty string to denote the start/end of a chain
_END_TOKEN_TEXT = ""
_NEXT_TOKEN_TABLE = "next_token"
_PREV_TOKEN_TABLE = "prev_token"

class Brain:
    def __init__(self, filename):
        self._db = db = Db(sqlite3.connect(filename))

        self.order = int(db.get_info_text("order"))

        self._end_token_id = db.get_token_id(_END_TOKEN_TEXT)

        self.tokenizer = tokenizer.MegaHALTokenizer()

    def learn(self, text):
        tokens = self.tokenizer.split(text.decode("utf-8"))

        if len(tokens) < self.order:
            log.debug("Input too short to learn: %s", text)
            return

        self._learn_tokens(tokens)

    def _learn_tokens(self, tokens):
        db = self._db
        c = db.cursor()

        token_ids = self._get_or_register_tokens(c, tokens)
        n_exprs = len(token_ids)-self.order

        for i in xrange(n_exprs+1):
            expr = token_ids[i:i+self.order]
            expr_id = self._get_or_register_expr(c, expr)

            # increment the expr count
            db.inc_expr_count(expr_id, c=c)

            if i == 0:
                # add link to boundary on prev_token
                db.add_or_inc_link(_PREV_TOKEN_TABLE, expr_id,
                                   self._end_token_id, c=c)

            if i > 0:
                # link prev token to this expr
                prev_token = token_ids[i-1]
                db.add_or_inc_link(_PREV_TOKEN_TABLE, expr_id, prev_token, c=c)

            if i < n_exprs:
                # link next token to this expr
                next_token = token_ids[i+self.order]
                db.add_or_inc_link(_NEXT_TOKEN_TABLE, expr_id, next_token, c=c)

            if i == n_exprs:
                # add link to boundary on next_token
                db.add_or_inc_link(_NEXT_TOKEN_TABLE, expr_id,
                                   self._end_token_id, c=c)

        db.commit()

    def reply(self, text):
        tokens = self.tokenizer.split(text.decode("utf-8"))

        db = self._db
        c = db.cursor()

        token_ids = self._get_known_tokens(tokens, c)

        best_score = 0.
        best_reply = None

        for i in xrange(100):
            (reply, score) = self._generate_reply(token_ids)

            if score > best_score:
                best_score = score
                best_reply = reply

        # look up the words for these tokens
        text = []
        for token_id in best_reply:
            text.append(db.get_token_text(token_id))

        return "".join(text)

    def _babble(self):
        return "I don't know enough to reply to that!"

    def _generate_reply(self, token_ids):
        # generate a reply containing one of token_ids

        db = self._db

        if len(token_ids) == 0:
            return self._babble()

        pivot_token_id = random.choice(token_ids)
        pivot_expr_id = db.get_random_expr(pivot_token_id)
        while pivot_expr_id is None:
            pivot_token_id = random.choice(token_ids)
            pivot_expr_id = db.get_random_expr(pivot_token_id)

        next_token_ids = db.follow_chain(_NEXT_TOKEN_TABLE, pivot_expr_id)
        prev_token_ids = db.follow_chain(_PREV_TOKEN_TABLE, pivot_expr_id)
        prev_token_ids.reverse()

        # strip the original expr from the prev reply
        prev_token_ids = prev_token_ids[:-self.order]

        all_tokens = prev_token_ids
        all_tokens.extend(next_token_ids)

        reply = []
        score = 0.

        for token in all_tokens:
            reply.append(token[0])
            score = score - math.log(token[1], 2)

        return reply, score

    def _get_known_tokens(self, tokens, c):
        db = self._db

        token_ids = []

        for token in tokens:
            token_id = db.get_token_id(token, c=c)
            if token_id is not None:
                token_ids.append(token_id)

        return token_ids

    def _get_or_register_tokens(self, c, tokens):
        db = self._db

        token_ids = []
        for token in tokens:
            token_id = db.get_token_id(token, c=c)
            if token_id is None:
                if re.search("\s", token):
                    is_whitespace = True
                else:
                    is_whitespace = False

                token_id = db.insert_token(token, is_whitespace, c=c)

            token_ids.append(token_id)

        return token_ids

    def _get_or_register_expr(self, c, token_ids):
        where = ["token%d_id = ?" % i for i in xrange(len(token_ids))]
        where = " AND ".join(where)

        row = c.execute("SELECT id FROM expr WHERE %s" % where, token_ids).fetchone()
        if row is None:
            cols = ",".join(["token%d_id" % i for i in xrange(len(token_ids))])
            params = ",".join(["?" for i in xrange(len(token_ids))])
            row = c.execute("INSERT INTO expr (count,%s) VALUES (0,%s)" %
                            (cols, params), token_ids)
            return c.lastrowid
        else:
            return row[0]

    @staticmethod
    def init(filename, order=5):
        """Initialize a brain. This brain's file must not already exist."""

        log.info("Initializing a hal brain.")

        db = Db(sqlite3.connect(filename))

        c = db.cursor()

        log.debug("Creating table: info")
        c.execute("""
CREATE TABLE info (
    attribute TEXT NOT NULL PRIMARY KEY,
    text TEXT NOT NULL)""")

        log.debug("Creating table: tokens")
        c.execute("""
CREATE TABLE tokens (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    text TEXT NOT NULL,
    is_whitespace INTEGER NOT NULL)""")

        tokens = []
        for i in xrange(order):
            tokens.append("token%d_id INTEGER NOT NULL REFERENCES token(id)" % i)

        log.debug("Creating table: expr")
        c.execute("""
CREATE TABLE expr (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    count INTEGER NOT NULL,
    %s)""" % ',\n    '.join(tokens))

        log.debug("Creating table: next_token")
        c.execute("""
CREATE TABLE next_token (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    expr_id INTEGER NOT NULL REFERENCES expr (id),
    token_id INTEGER NOT NULL REFERENCES token (id),
    count INTEGER NOT NULL)""")

        log.debug("Creating table: prev_token")
        c.execute("""
CREATE TABLE prev_token (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    expr_id INTEGER NOT NULL REFERENCES expr (id),
    token_id INTEGER NOT NULL REFERENCES token (id),
    count INTEGER NOT NULL)""")

        # create a token for the end of a chain
        db.insert_token(_END_TOKEN_TEXT, 0, c=c)

        # save the order of this brain
        db.set_info_text("order", str(order), c=c)

        db.commit()
        c.close()
        db.close()

    def clone(self, megahal_brain):
        log.info("Cloning a MegaHAL brain: %s", megahal_brain)

        conn = sqlite3.connect(self.filename)

        b = megahal.Brain(megahal_brain)
        b.clone(conn)

        conn.close()

class Db:
    """Database functions to support a Hal brain."""
    def __init__(self, conn):
        self._conn = conn

        if self.is_initted():
            self._order = int(self.get_info_text("order"))
            self._end_token_id = self.get_token_id(_END_TOKEN_TEXT)

            self._all_tokens = ",".join(["token%d_id" % i
                                         for i in xrange(self._order)])
            self._all_token_args = " AND ".join(["token%d_id = ?" % i
                                                 for i in xrange(self._order)])

    def cursor(self):
        return self._conn.cursor()

    def commit(self):
        return self._conn.commit()

    def close(self):
        return self._conn.close()

    def is_initted(self, c=None):
        if c is None:
            c = self.cursor()

        try:
            self.get_info_text("order")
            return True
        except sqlite3.OperationalError:
            return False

    def set_info_text(self, attribute, text, c=None):
        if c is None:
            c = self.cursor()

        # FIXME: needs to UPDATE if the key already exists

        q = "INSERT INTO info (attribute, text) VALUES (?, ?)"
        c.execute(q, (attribute, text))

    def get_info_text(self, attribute, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT text FROM info WHERE attribute = ?"
        row = c.execute(q, (attribute,)).fetchone()
        if row:
            return row[0]

    def get_token_id(self, token, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT id FROM tokens WHERE text = ?"
        row = c.execute(q, (token,)).fetchone()
        if row:
            return int(row[0])

    def get_token_text(self, token_id, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT text FROM tokens WHERE id = ?"
        row = c.execute(q, (token_id,)).fetchone()
        if row:
            return row[0]

    def get_token_texts(self, token_ids, c=None):
        if c is None:
            c = self.cursor()

        return "".join([self.get_token_text(token_id)
                        for token_id in token_ids])

    def _get_expr_token_ids(self, expr_id, c):
        q = "SELECT count,%s FROM expr WHERE id = ?" % self._all_tokens
        return c.execute(q, (expr_id,)).fetchone()

    def insert_token(self, token, is_whitespace, c=None):
        if c is None:
            c = self.cursor()

        q = "INSERT INTO tokens (text, is_whitespace) VALUES (?, ?)"
        c.execute(q, (token, is_whitespace))
        return c.lastrowid

    def inc_expr_count(self, expr_id, c=None):
        if c is None:
            c = self.cursor()

        q = "UPDATE expr SET count = count + 1 WHERE id = ?"
        c.execute(q, (expr_id,))

    def _count_expr_token_links(self, table, expr_id, token_id, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT count FROM %s WHERE expr_id = ? AND token_id = ?" % table
        row = c.execute(q, (expr_id, token_id)).fetchone()
        if row:
            return int(row[0])

    def add_or_inc_link(self, table, expr_id, token_id, c=None):
        if c is None:
            c = self.cursor()

        count = self._count_expr_token_links(table, expr_id, token_id, c)

        if count is not None:
            q = "UPDATE %s SET count = count + 1 WHERE expr_id = ? AND token_id = ?" % table
            c.execute(q, (expr_id, token_id))
        else:
            q = "INSERT INTO %s (expr_id, token_id, count) VALUES (?, ?, ?)" % table
            c.execute(q, (expr_id, token_id, 1))

    def get_random_expr(self, token_id, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT id FROM expr WHERE token0_id = ? ORDER BY RANDOM()"
        row = c.execute(q, (token_id,)).fetchone()
        if row:
            return int(row[0])

    def _get_expr_by_token_ids(self, token_ids, c):
        q = "SELECT id, count FROM expr WHERE %s" % self._all_token_args

        row = c.execute(q, token_ids).fetchone()
        if row:
            return int(row[0]), int(row[1])

    def _get_expr_count(self, expr_id, c):
        q = "SELECT count FROM expr WHERE id = ?"

        row = c.execute(q, expr_id).fetchone()
        if row:
            return int(row[0])

    def follow_chain(self, table, expr_id, c=None):
        if c is None:
            c = self.cursor()

        expr_info = self._get_expr_token_ids(expr_id, c)
        next_expr_count = expr_info[0]
        next_token_ids = expr_info[1:]

        chain = [(token_id, 1.0) for token_id in next_token_ids]

        # pick a random next_token given the things in expr_id
        q = "SELECT token_id, count FROM %s WHERE expr_id = ? ORDER BY RANDOM()" % table

        row = c.execute(q, (expr_id,)).fetchone()
        next_token_id, next_token_count = row

        while next_token_id != self._end_token_id:
            chain.append((next_token_id,
                          float(next_token_count) / next_expr_count))

            if table == _NEXT_TOKEN_TABLE:
                next_token_ids = list(next_token_ids[1:])
                next_token_ids.append(next_token_id)
            else:
                new_next_token_ids = [next_token_id]
                new_next_token_ids.extend(next_token_ids[:-1])
                next_token_ids = new_next_token_ids

            next_expr_id, next_expr_count = self._get_expr_by_token_ids(next_token_ids, c)

            row = c.execute(q, (next_expr_id,)).fetchone()
            next_token_id, next_token_count = row

        return chain
<commit_msg>Add database indexes.<commit_after># Copyright (C) 2010 Peter Teichman

import logging
import math
import random
import re
import sqlite3

import megahal
import tokenizer

log = logging.getLogger("hal")

# use an empty string to denote the start/end of a chain
_END_TOKEN_TEXT = ""
_NEXT_TOKEN_TABLE = "next_token"
_PREV_TOKEN_TABLE = "prev_token"

class Brain:
    def __init__(self, filename):
        self._db = db = Db(sqlite3.connect(filename))

        self.order = int(db.get_info_text("order"))

        self._end_token_id = db.get_token_id(_END_TOKEN_TEXT)

        self.tokenizer = tokenizer.MegaHALTokenizer()

    def learn(self, text):
        tokens = self.tokenizer.split(text.decode("utf-8"))

        if len(tokens) < self.order:
            log.debug("Input too short to learn: %s", text)
            return

        self._learn_tokens(tokens)

    def _learn_tokens(self, tokens):
        db = self._db
        c = db.cursor()

        token_ids = self._get_or_register_tokens(c, tokens)
        n_exprs = len(token_ids)-self.order

        for i in xrange(n_exprs+1):
            expr = token_ids[i:i+self.order]
            expr_id = self._get_or_register_expr(c, expr)

            # increment the expr count
            db.inc_expr_count(expr_id, c=c)

            if i == 0:
                # add link to boundary on prev_token
                db.add_or_inc_link(_PREV_TOKEN_TABLE, expr_id,
                                   self._end_token_id, c=c)

            if i > 0:
                # link prev token to this expr
                prev_token = token_ids[i-1]
                db.add_or_inc_link(_PREV_TOKEN_TABLE, expr_id, prev_token, c=c)

            if i < n_exprs:
                # link next token to this expr
                next_token = token_ids[i+self.order]
                db.add_or_inc_link(_NEXT_TOKEN_TABLE, expr_id, next_token, c=c)

            if i == n_exprs:
                # add link to boundary on next_token
                db.add_or_inc_link(_NEXT_TOKEN_TABLE, expr_id,
                                   self._end_token_id, c=c)

        db.commit()

    def reply(self, text):
        tokens = self.tokenizer.split(text.decode("utf-8"))

        db = self._db
        c = db.cursor()

        token_ids = self._get_known_tokens(tokens, c)

        best_score = 0.
        best_reply = None

        for i in xrange(100):
            (reply, score) = self._generate_reply(token_ids)

            if score > best_score:
                best_score = score
                best_reply = reply

        if best_reply is None:
            return "I don't know enough to answer you yet!"

        # look up the words for these tokens
        text = []
        for token_id in best_reply:
            text.append(db.get_token_text(token_id))

        return "".join(text)

    def _babble(self):
        return "I don't know enough to reply to that!"

    def _generate_reply(self, token_ids):
        # generate a reply containing one of token_ids

        db = self._db

        if len(token_ids) == 0:
            return self._babble()

        pivot_token_id = random.choice(token_ids)
        pivot_expr_id = db.get_random_expr(pivot_token_id)
        while pivot_expr_id is None:
            pivot_token_id = random.choice(token_ids)
            pivot_expr_id = db.get_random_expr(pivot_token_id)

        next_token_ids = db.follow_chain(_NEXT_TOKEN_TABLE, pivot_expr_id)
        prev_token_ids = db.follow_chain(_PREV_TOKEN_TABLE, pivot_expr_id)
        prev_token_ids.reverse()

        # strip the original expr from the prev reply
        prev_token_ids = prev_token_ids[:-self.order]

        all_tokens = prev_token_ids
        all_tokens.extend(next_token_ids)

        reply = []
        score = 0.

        for token in all_tokens:
            reply.append(token[0])
            score = score - math.log(token[1], 2)

        return reply, score

    def _get_known_tokens(self, tokens, c):
        db = self._db

        token_ids = []

        for token in tokens:
            token_id = db.get_token_id(token, c=c)
            if token_id is not None:
                token_ids.append(token_id)

        return token_ids

    def _get_or_register_tokens(self, c, tokens):
        db = self._db

        token_ids = []
        for token in tokens:
            token_id = db.get_token_id(token, c=c)
            if token_id is None:
                if re.search("\s", token):
                    is_whitespace = True
                else:
                    is_whitespace = False

                token_id = db.insert_token(token, is_whitespace, c=c)

            token_ids.append(token_id)

        return token_ids

    def _get_or_register_expr(self, c, token_ids):
        where = ["token%d_id = ?" % i for i in xrange(len(token_ids))]
        where = " AND ".join(where)

        row = c.execute("SELECT id FROM expr WHERE %s" % where, token_ids).fetchone()
        if row is None:
            cols = ",".join(["token%d_id" % i for i in xrange(len(token_ids))])
            params = ",".join(["?" for i in xrange(len(token_ids))])
            row = c.execute("INSERT INTO expr (count,%s) VALUES (0,%s)" %
                            (cols, params), token_ids)
            return c.lastrowid
        else:
            return row[0]

    @staticmethod
    def init(filename, order=5, create_indexes=True):
        """Initialize a brain. This brain's file must not already exist."""

        log.info("Initializing a hal brain.")

        db = Db(sqlite3.connect(filename))

        c = db.cursor()

        log.debug("Creating table: info")
        c.execute("""
CREATE TABLE info (
    attribute TEXT NOT NULL PRIMARY KEY,
    text TEXT NOT NULL)""")

        log.debug("Creating table: tokens")
        c.execute("""
CREATE TABLE tokens (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    text TEXT NOT NULL,
    is_whitespace INTEGER NOT NULL)""")

        tokens = []
        for i in xrange(order):
            tokens.append("token%d_id INTEGER NOT NULL REFERENCES token(id)" % i)

        log.debug("Creating table: expr")
        c.execute("""
CREATE TABLE expr (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    count INTEGER NOT NULL,
    %s)""" % ',\n    '.join(tokens))

        log.debug("Creating table: next_token")
        c.execute("""
CREATE TABLE next_token (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    expr_id INTEGER NOT NULL REFERENCES expr (id),
    token_id INTEGER NOT NULL REFERENCES token (id),
    count INTEGER NOT NULL)""")

        log.debug("Creating table: prev_token")
        c.execute("""
CREATE TABLE prev_token (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    expr_id INTEGER NOT NULL REFERENCES expr (id),
    token_id INTEGER NOT NULL REFERENCES token (id),
    count INTEGER NOT NULL)""")

        # create a token for the end of a chain
        db.insert_token(_END_TOKEN_TEXT, 0, c=c)

        # save the order of this brain
        db.set_info_text("order", str(order), c=c)

        if create_indexes:
            c.execute("""
CREATE INDEX tokens_text on tokens (text)""")

            for i in xrange(order):
                c.execute("""
CREATE INDEX expr_token%d_id on expr (token%d_id)""" % (i, i))

            token_ids = ",".join(["token%d_id" % i for i in xrange(order)])
            c.execute("""
CREATE INDEX expr_token_ids on expr (%s)""" % token_ids)

            c.execute("""
CREATE INDEX next_token_expr_id ON next_token (expr_id)""")

            c.execute("""
CREATE INDEX prev_token_expr_id ON prev_token (expr_id)""")

        db.commit()
        c.close()
        db.close()

    def clone(self, megahal_brain):
        log.info("Cloning a MegaHAL brain: %s", megahal_brain)

        conn = sqlite3.connect(self.filename)

        b = megahal.Brain(megahal_brain)
        b.clone(conn)

        conn.close()

class Db:
    """Database functions to support a Hal brain."""
    def __init__(self, conn):
        self._conn = conn

        if self.is_initted():
            self._order = int(self.get_info_text("order"))
            self._end_token_id = self.get_token_id(_END_TOKEN_TEXT)

            self._all_tokens = ",".join(["token%d_id" % i
                                         for i in xrange(self._order)])
            self._all_token_args = " AND ".join(["token%d_id = ?" % i
                                                 for i in xrange(self._order)])

    def cursor(self):
        return self._conn.cursor()

    def commit(self):
        return self._conn.commit()

    def close(self):
        return self._conn.close()

    def is_initted(self, c=None):
        if c is None:
            c = self.cursor()

        try:
            self.get_info_text("order")
            return True
        except sqlite3.OperationalError:
            return False

    def set_info_text(self, attribute, text, c=None):
        if c is None:
            c = self.cursor()

        # FIXME: needs to UPDATE if the key already exists

        q = "INSERT INTO info (attribute, text) VALUES (?, ?)"
        c.execute(q, (attribute, text))

    def get_info_text(self, attribute, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT text FROM info WHERE attribute = ?"
        row = c.execute(q, (attribute,)).fetchone()
        if row:
            return row[0]

    def get_token_id(self, token, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT id FROM tokens WHERE text = ?"
        row = c.execute(q, (token,)).fetchone()
        if row:
            return int(row[0])

    def get_token_text(self, token_id, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT text FROM tokens WHERE id = ?"
        row = c.execute(q, (token_id,)).fetchone()
        if row:
            return row[0]

    def get_token_texts(self, token_ids, c=None):
        if c is None:
            c = self.cursor()

        return "".join([self.get_token_text(token_id)
                        for token_id in token_ids])

    def _get_expr_token_ids(self, expr_id, c):
        q = "SELECT count,%s FROM expr WHERE id = ?" % self._all_tokens
        return c.execute(q, (expr_id,)).fetchone()

    def insert_token(self, token, is_whitespace, c=None):
        if c is None:
            c = self.cursor()

        q = "INSERT INTO tokens (text, is_whitespace) VALUES (?, ?)"
        c.execute(q, (token, is_whitespace))
        return c.lastrowid

    def inc_expr_count(self, expr_id, c=None):
        if c is None:
            c = self.cursor()

        q = "UPDATE expr SET count = count + 1 WHERE id = ?"
        c.execute(q, (expr_id,))

    def _count_expr_token_links(self, table, expr_id, token_id, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT count FROM %s WHERE expr_id = ? AND token_id = ?" % table
        row = c.execute(q, (expr_id, token_id)).fetchone()
        if row:
            return int(row[0])

    def add_or_inc_link(self, table, expr_id, token_id, c=None):
        if c is None:
            c = self.cursor()

        count = self._count_expr_token_links(table, expr_id, token_id, c)

        if count is not None:
            q = "UPDATE %s SET count = count + 1 WHERE expr_id = ? AND token_id = ?" % table
            c.execute(q, (expr_id, token_id))
        else:
            q = "INSERT INTO %s (expr_id, token_id, count) VALUES (?, ?, ?)" % table
            c.execute(q, (expr_id, token_id, 1))

    def get_random_expr(self, token_id, c=None):
        if c is None:
            c = self.cursor()

        q = "SELECT id FROM expr WHERE token0_id = ? ORDER BY RANDOM()"
        row = c.execute(q, (token_id,)).fetchone()
        if row:
            return int(row[0])

    def _get_expr_by_token_ids(self, token_ids, c):
        q = "SELECT id, count FROM expr WHERE %s" % self._all_token_args

        row = c.execute(q, token_ids).fetchone()
        if row:
            return int(row[0]), int(row[1])

    def _get_expr_count(self, expr_id, c):
        q = "SELECT count FROM expr WHERE id = ?"

        row = c.execute(q, expr_id).fetchone()
        if row:
            return int(row[0])

    def follow_chain(self, table, expr_id, c=None):
        if c is None:
            c = self.cursor()

        expr_info = self._get_expr_token_ids(expr_id, c)
        next_expr_count = expr_info[0]
        next_token_ids = expr_info[1:]

        chain = [(token_id, 1.0) for token_id in next_token_ids]

        # pick a random next_token given the things in expr_id
        q = "SELECT token_id, count FROM %s WHERE expr_id = ? ORDER BY RANDOM()" % table

        row = c.execute(q, (expr_id,)).fetchone()
        next_token_id, next_token_count = row

        while next_token_id != self._end_token_id:
            chain.append((next_token_id,
                          float(next_token_count) / next_expr_count))

            if table == _NEXT_TOKEN_TABLE:
                next_token_ids = list(next_token_ids[1:])
                next_token_ids.append(next_token_id)
            else:
                new_next_token_ids = [next_token_id]
                new_next_token_ids.extend(next_token_ids[:-1])
                next_token_ids = new_next_token_ids

            next_expr_id, next_expr_count = self._get_expr_by_token_ids(next_token_ids, c)

            row = c.execute(q, (next_expr_id,)).fetchone()
            next_token_id, next_token_count = row

        return chain
<|endoftext|>


<commit_before>"""

Copyright 2015, Institute for Systems Biology

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.

"""

import logging
from re import compile as re_compile

from bq_data_access.v2.errors import FeatureNotFoundException
from bq_data_access.v2.feature_value_types import ValueType, DataTypes
from bq_data_access.v2.utils import DurationLogged
from bq_data_access.data_types.gnab import BIGQUERY_CONFIG
from scripts.feature_def_gen.gnab_features import GNABFeatureDefConfig

GNAB_FEATURE_TYPE = 'GNAB'
IDENTIFIER_COLUMN_NAME = 'sample_id'

def get_feature_type():
    return GNAB_FEATURE_TYPE

class GNABFeatureDef(object):
    VALUE_FIELD_NUM_MUTATIONS = 'num_mutations'
    config_instance = GNABFeatureDefConfig.from_dict(BIGQUERY_CONFIG)

    # Regular expression for parsing the feature definition.
    #
    # Example ID: v2:GNAB:SMYD3:hg19_mc3:Variant_Classification
    regex = re_compile("^v2:GNAB:"
                       # gene
                       "([a-zA-Z0-9_.\-]+):"
                       "(" + "|".join([table.internal_table_id for table in config_instance.data_table_list]) + "):"
                       # value field
                       "(Variant_Classification|Variant_Type|{})$".format(VALUE_FIELD_NUM_MUTATIONS))

    def __init__(self, gene, internal_table_id, value_field):
        self.gene = gene
        self.internal_table_id = internal_table_id
        self.value_field = value_field

    def get_table_configuration(self):
        for table_config in self.config_instance.data_table_list:
            if table_config.internal_table_id == self.internal_table_id:
                return table_config

    @classmethod
    def from_feature_id(cls, feature_id):
        feature_fields = cls.regex.findall(feature_id)
        if len(feature_fields) == 0:
            raise FeatureNotFoundException(feature_id)

        gene_label, internal_table_id, value_field = feature_fields[0]

        return cls(gene_label, internal_table_id, value_field)

class GNABFeatureProvider(object):
    def __init__(self, feature_id, **kwargs):
        self.feature_def = None
        self.table_info = None
        self.table_name = ''
        self.parse_internal_feature_id(feature_id)
        self.config_instance = GNABFeatureDefConfig.from_dict(BIGQUERY_CONFIG)

    def get_value_type(self):
        if self.feature_def.value_field == GNABFeatureDef.VALUE_FIELD_NUM_MUTATIONS:
            return ValueType.FLOAT
        else:
            return ValueType.STRING

    def get_feature_type(self):
        return DataTypes.GNAB

    @classmethod
    def process_data_point(cls, data_point):
        return data_point['value']

    def build_query_for_program(self, feature_def, cohort_table, cohort_id_array, project_id_array):
        # Generate the 'IN' statement string: (%s, %s, ..., %s)
        cohort_id_stmt = ', '.join([str(cohort_id) for cohort_id in cohort_id_array])
        project_id_stmt = ''
        if project_id_array is not None:
            project_id_stmt = ', '.join([str(project_id) for project_id in project_id_array])

        query_template = "SELECT case_barcode, sample_barcode_tumor, aliquot_barcode_tumor, " \
             "{value_field} AS value {brk}" \
             "FROM [{table_name}] AS gnab {brk}" \
             "WHERE {gene_label_field}='{gene_symbol}' {brk}" \
             "AND sample_barcode_tumor IN ( {brk}" \
             "    SELECT sample_barcode {brk}" \
             "    FROM [{cohort_dataset_and_table}] {brk}" \
             "    WHERE cohort_id IN ({cohort_id_list}) {brk}" \
             "         AND (project_id IS NULL {brk}"

        query_template += (" OR project_id IN ({project_id_list}))) {brk}" if project_id_array is not None else ")) {brk}")

        value_field_bqsql = self.feature_def.value_field

        if self.feature_def.value_field == GNABFeatureDef.VALUE_FIELD_NUM_MUTATIONS:
            value_field_bqsql = 'count(*)'
            query_template += ("GROUP BY case_barcode, sample_barcode_tumor, aliquot_barcode_tumor {brk}")

        table_config = feature_def.get_table_configuration()

        query = query_template.format(table_name=table_config.table_id,
                                      gene_label_field=table_config.gene_label_field,
                                      gene_symbol=feature_def.gene,
                                      value_field=value_field_bqsql,
                                      cohort_dataset_and_table=cohort_table,
                                      cohort_id_list=cohort_id_stmt, project_id_list=project_id_stmt,
                                      brk='\n')

        logging.debug("BQ_QUERY_GNAB: " + query)
        return query

    def build_query(self, project_set, cohort_table, cohort_id_array, project_id_array):
        query = self.build_query_for_program(self.feature_def, cohort_table, cohort_id_array, project_id_array)
        return query

    @DurationLogged('GNAB', 'UNPACK')
    def unpack_query_response(self, query_result_array):
        """
        Unpacks values from a BigQuery response object into a flat array. The array will contain dicts with
        the following fields:
        - 'patient_id': Patient barcode
        - 'sample_id': Sample barcode
        - 'aliquot_id': Aliquot barcode
        - 'value': Value of the selected column from the MAF data table

        Args:
            query_result_array: A BigQuery query response object

        Returns:
            Array of dict objects.
        """
        result = []

        for row in query_result_array:
            result.append({
                'case_id': row['f'][0]['v'],
                'sample_id': row['f'][1]['v'],
                'aliquot_id': row['f'][2]['v'],
                'value': row['f'][3]['v'],
            })

        return result

    def parse_internal_feature_id(self, feature_id):
        self.feature_def = GNABFeatureDef.from_feature_id(feature_id)

    @classmethod
    def is_valid_feature_id(cls, feature_id):
        is_valid = False
        try:
            GNABFeatureDef.from_feature_id(feature_id)
            is_valid = True
        except Exception:
            # GNABFeatureDef.from_feature_id raises Exception if the feature identifier
            # is not valid. Nothing needs to be done here, since is_valid is already False.
            pass
        finally:
            return is_valid
<commit_msg>Cleanup<commit_after>"""

Copyright 2015, Institute for Systems Biology

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.

"""

import logging
from re import compile as re_compile

from bq_data_access.v2.errors import FeatureNotFoundException
from bq_data_access.v2.feature_value_types import ValueType, DataTypes
from bq_data_access.v2.utils import DurationLogged
from bq_data_access.data_types.gnab import BIGQUERY_CONFIG
from scripts.feature_def_gen.gnab_features import GNABFeatureDefConfig

GNAB_FEATURE_TYPE = 'GNAB'
IDENTIFIER_COLUMN_NAME = 'sample_id'

def get_feature_type():
    return GNAB_FEATURE_TYPE

class GNABFeatureDef(object):
    VALUE_FIELD_NUM_MUTATIONS = 'num_mutations'
    config_instance = GNABFeatureDefConfig.from_dict(BIGQUERY_CONFIG)

    # Regular expression for parsing the feature definition.
    #
    # Example ID: v2:GNAB:SMYD3:hg19_mc3:Variant_Classification
    regex = re_compile("^v2:GNAB:"
                       # gene
                       "([a-zA-Z0-9_.\-]+):"
                       "(" + "|".join([table.internal_table_id for table in config_instance.data_table_list]) + "):"
                       # value field
                       "(Variant_Classification|Variant_Type|{})$".format(VALUE_FIELD_NUM_MUTATIONS))

    def __init__(self, gene, internal_table_id, value_field):
        self.gene = gene
        self.internal_table_id = internal_table_id
        self.value_field = value_field

    def get_table_configuration(self):
        for table_config in self.config_instance.data_table_list:
            if table_config.internal_table_id == self.internal_table_id:
                return table_config

    @classmethod
    def from_feature_id(cls, feature_id):
        feature_fields = cls.regex.findall(feature_id)
        if len(feature_fields) == 0:
            raise FeatureNotFoundException(feature_id)

        gene_label, internal_table_id, value_field = feature_fields[0]

        return cls(gene_label, internal_table_id, value_field)

class GNABFeatureProvider(object):
    def __init__(self, feature_id, **kwargs):
        self.feature_def = None
        self.parse_internal_feature_id(feature_id)
        self.config_instance = GNABFeatureDefConfig.from_dict(BIGQUERY_CONFIG)

    def get_value_type(self):
        if self.feature_def.value_field == GNABFeatureDef.VALUE_FIELD_NUM_MUTATIONS:
            return ValueType.FLOAT
        else:
            return ValueType.STRING

    def get_feature_type(self):
        return DataTypes.GNAB

    @classmethod
    def process_data_point(cls, data_point):
        return data_point['value']

    def build_query_for_program(self, feature_def, cohort_table, cohort_id_array, project_id_array):
        # Generate the 'IN' statement string: (%s, %s, ..., %s)
        cohort_id_stmt = ', '.join([str(cohort_id) for cohort_id in cohort_id_array])
        project_id_stmt = ''
        if project_id_array is not None:
            project_id_stmt = ', '.join([str(project_id) for project_id in project_id_array])

        query_template = "SELECT case_barcode, sample_barcode_tumor, aliquot_barcode_tumor, " \
             "{value_field} AS value {brk}" \
             "FROM [{table_name}] AS gnab {brk}" \
             "WHERE {gene_label_field}='{gene_symbol}' {brk}" \
             "AND sample_barcode_tumor IN ( {brk}" \
             "    SELECT sample_barcode {brk}" \
             "    FROM [{cohort_dataset_and_table}] {brk}" \
             "    WHERE cohort_id IN ({cohort_id_list}) {brk}" \
             "         AND (project_id IS NULL {brk}"

        query_template += (" OR project_id IN ({project_id_list}))) {brk}" if project_id_array is not None else ")) {brk}")

        value_field_bqsql = self.feature_def.value_field

        if self.feature_def.value_field == GNABFeatureDef.VALUE_FIELD_NUM_MUTATIONS:
            value_field_bqsql = 'count(*)'
            query_template += ("GROUP BY case_barcode, sample_barcode_tumor, aliquot_barcode_tumor {brk}")

        table_config = feature_def.get_table_configuration()

        query = query_template.format(table_name=table_config.table_id,
                                      gene_label_field=table_config.gene_label_field,
                                      gene_symbol=feature_def.gene,
                                      value_field=value_field_bqsql,
                                      cohort_dataset_and_table=cohort_table,
                                      cohort_id_list=cohort_id_stmt, project_id_list=project_id_stmt,
                                      brk='\n')

        logging.debug("BQ_QUERY_GNAB: " + query)
        return query

    def build_query(self, project_set, cohort_table, cohort_id_array, project_id_array):
        query = self.build_query_for_program(self.feature_def, cohort_table, cohort_id_array, project_id_array)
        return query

    @DurationLogged('GNAB', 'UNPACK')
    def unpack_query_response(self, query_result_array):
        """
        Unpacks values from a BigQuery response object into a flat array. The array will contain dicts with
        the following fields:
        - 'patient_id': Patient barcode
        - 'sample_id': Sample barcode
        - 'aliquot_id': Aliquot barcode
        - 'value': Value of the selected column from the MAF data table

        Args:
            query_result_array: A BigQuery query response object

        Returns:
            Array of dict objects.
        """
        result = []

        for row in query_result_array:
            result.append({
                'case_id': row['f'][0]['v'],
                'sample_id': row['f'][1]['v'],
                'aliquot_id': row['f'][2]['v'],
                'value': row['f'][3]['v'],
            })

        return result

    def parse_internal_feature_id(self, feature_id):
        self.feature_def = GNABFeatureDef.from_feature_id(feature_id)

    @classmethod
    def is_valid_feature_id(cls, feature_id):
        is_valid = False
        try:
            GNABFeatureDef.from_feature_id(feature_id)
            is_valid = True
        except Exception:
            # GNABFeatureDef.from_feature_id raises Exception if the feature identifier
            # is not valid. Nothing needs to be done here, since is_valid is already False.
            pass
        finally:
            return is_valid
<|endoftext|>