<commit_before>"""
Mask specified sites from a VCF or FASTA file.
"""
import os
import sys
from shutil import copyfile

import numpy as np
import pandas as pd
from Bio import SeqIO
from Bio.Seq import MutableSeq

from .utils import run_shell_command, shquote, open_file, is_vcf

def get_chrom_name(vcf_file):
    """Read the CHROM field from the first non-header line of a vcf file.

    Returns:
    str or None: Either the CHROM field or None if no non-comment line could be found.
    """
    with open_file(vcf_file, mode='r') as f:
        for line in f:
            if line[0] != "#":
                header = line.strip().partition('\t')
                return header[0]

def read_bed_file(mask_file):
    """Read the full list of excluded sites from the BED file.

    Second column is chromStart, 3rd is chromEnd. Generate a range from these two columns.
    """
    sites_to_mask = []
    bed = pd.read_csv(mask_file, sep='\t', header=None, usecols=[1,2])
    for idx, row in bed.iterrows():
        try:
            sites_to_mask.extend(list(range(int(row[1]), int(row[2])+1)))
        except ValueError as err:
            # Skip unparseable lines, including header lines.
            print("Could not read line %d of BED file %s: %s. Continuing." % (idx, mask_file, err))
    sites_to_mask = np.unique(sites_to_mask).tolist()
    print("Found %d sites to mask" % len(sites_to_mask))
    return sites_to_mask

def mask_vcf(mask_sites, in_file, out_file, cleanup=True):
    """Mask the provided site list from a VCF file and write to a new file.

    This function relies on 'vcftools --exclude-positions' to mask the requested sites.

    Parameters:
    -----------
    mask_sites: list[int]
        A list of site indexes to exclude from the vcf.
    in_file: str
        The path to the vcf file you wish to mask.
    out_file: str
        The path to write the resulting vcf to
    cleanup: bool
        Clean up the intermediate files, including the VCFTools log and mask sites file
    """
    cleanup_files = ['out.log']

    # Create the temporary masking file to be passed to VCFTools
    # Need CHROM name from VCF file:
    chrom_name = get_chrom_name(in_file)
    if chrom_name is None:
        print("ERROR: Something went wrong reading your VCF file: a CHROM column could not be found. "
              "Please check the file is valid VCF format.")
        sys.exit(1)

    exclude = [chrom_name + "\t" + str(pos) for pos in mask_sites]
    temp_mask_file = in_file + "_maskTemp"
    with open_file(temp_mask_file, 'w') as fh:
        fh.write("\n".join(exclude))
    cleanup_files.append(temp_mask_file)

    #Read in/write out according to file ending
    in_call = "--gzvcf" if in_file.lower().endswith(".gz") else "--vcf"
    out_call = "| gzip -c" if out_file.lower().endswith(".gz") else ""

    call = ["vcftools", "--exclude-positions", shquote(temp_mask_file), in_call, shquote(in_file), "--recode --stdout", out_call, ">", shquote(out_file)]
    print("Removing masked sites from VCF file using vcftools... this may take some time. Call:")
    print(" ".join(call))
    run_shell_command(" ".join(call), raise_errors = True)
    # remove vcftools log file
    if cleanup:
        for file in cleanup_files:
            try:
                os.remove(file)
            except OSError:
                pass

def mask_fasta(mask_sites, in_file, out_file, mask_from_beginning=0, mask_from_end=0):
    """Mask the provided site list from a FASTA file and write to a new file.

    Masked sites are overwritten as "N"s.

:
    -----------
    mask_sites: list[int]
        A list of site indexes to exclude from the FASTA.
    in_file: str
        The path to the FASTA file you wish to mask.
    out_file: str
        The path to write the resulting FASTA to
    mask_from_beginning: int
       Number of sites to mask from the beginning of each sequence (default 0)
    mask_from_end: int
       Number of sites to mask from the end of each sequence (default 0)
    """
    # Load alignment as FASTA generator to prevent loading the whole alignment
    # into memory.
    alignment = SeqIO.parse(in_file, "fasta")

    # Write the masked alignment to disk one record at a time.
    print("Removing masked sites from FASTA file.")
    with open_file(out_file, "w") as oh:
        for record in alignment:
            # Convert to a mutable sequence to enable masking with Ns.
            sequence = MutableSeq(
                "N" * mask_from_beginning +
                str(record.seq)[mask_from_beginning:-mask_from_end or None] +
                "N" * mask_from_end
            )
            sequence_length = len(sequence)
            # Replace all excluded sites with Ns.
            for site in mask_sites:
                if site < sequence_length:
                    sequence[site] = "N"
            record.seq = sequence
            SeqIO.write(record, oh, "fasta")

def register_arguments(parser):
    parser.add_argument('--sequences', '-s', required=True, help="sequences in VCF or FASTA format")
    parser.add_argument('--mask', required=True, help="locations to be masked in BED file format")
    parser.add_argument('--output', '-o', help="output file")
    parser.add_argument('--no-cleanup', dest="cleanup", action="store_false",
                        help="Leave intermediate files around. May be useful for debugging")

def run(args):
    '''
    Mask specified sites from the VCF or FASTA.

    For VCF files, his occurs by removing them entirely from the VCF, essentially making
    them identical to the reference at the locations.

    For FASTA files, masked sites are replaced with "N".

    If users don't specify output, will overwrite the input file.
    '''
    # Check files exist and are not empty
    if not os.path.isfile(args.sequences):
        print("ERROR: File {} does not exist!".format(args.sequences))
        return 1
    if not os.path.isfile(args.mask):
        print("ERROR: File {} does not exist!".format(args.mask))
        return 1
    if os.path.getsize(args.sequences) == 0:
        print("ERROR: {} is empty. Please check how this file was produced. "
              "Did an error occur in an earlier step?".format(args.sequences))
        return 1
    if os.path.getsize(args.mask) == 0:
        print("ERROR: {} is an empty file.".format(args.mask))
        return 1

    mask_sites = read_bed_file(args.mask)

    # For both FASTA and VCF masking, we need a proper separate output file
    if args.output is not None:
        out_file = args.output
    else:
        out_file = os.path.join(os.path.dirname(args.sequences),
                                "masked_" + os.path.basename(args.sequences))

    if is_vcf(args.sequences):
        mask_vcf(mask_sites, args.sequences, out_file, args.cleanup)
    else:
        mask_fasta(mask_sites, args.sequences, out_file)

    if args.output is None:
        copyfile(out_file, args.sequences)
        if args.cleanup:
            os.remove(out_file)
<commit_msg>Make args.mask optional, change dest to args.mask_file<commit_after>"""
Mask specified sites from a VCF or FASTA file.
"""
import os
import sys
from shutil import copyfile

import numpy as np
import pandas as pd
from Bio import SeqIO
from Bio.Seq import MutableSeq

from .utils import run_shell_command, shquote, open_file, is_vcf

def get_chrom_name(vcf_file):
    """Read the CHROM field from the first non-header line of a vcf file.

    Returns:
    str or None: Either the CHROM field or None if no non-comment line could be found.
    """
    with open_file(vcf_file, mode='r') as f:
        for line in f:
            if line[0] != "#":
                header = line.strip().partition('\t')
                return header[0]

def read_bed_file(mask_file):
    """Read the full list of excluded sites from the BED file.

    Second column is chromStart, 3rd is chromEnd. Generate a range from these two columns.
    """
    sites_to_mask = []
    bed = pd.read_csv(mask_file, sep='\t', header=None, usecols=[1,2])
    for idx, row in bed.iterrows():
        try:
            sites_to_mask.extend(list(range(int(row[1]), int(row[2])+1)))
        except ValueError as err:
            # Skip unparseable lines, including header lines.
            print("Could not read line %d of BED file %s: %s. Continuing." % (idx, mask_file, err))
    sites_to_mask = np.unique(sites_to_mask).tolist()
    print("Found %d sites to mask" % len(sites_to_mask))
    return sites_to_mask

def mask_vcf(mask_sites, in_file, out_file, cleanup=True):
    """Mask the provided site list from a VCF file and write to a new file.

    This function relies on 'vcftools --exclude-positions' to mask the requested sites.

    Parameters:
    -----------
    mask_sites: list[int]
        A list of site indexes to exclude from the vcf.
    in_file: str
        The path to the vcf file you wish to mask.
    out_file: str
        The path to write the resulting vcf to
    cleanup: bool
        Clean up the intermediate files, including the VCFTools log and mask sites file
    """
    cleanup_files = ['out.log']

    # Create the temporary masking file to be passed to VCFTools
    # Need CHROM name from VCF file:
    chrom_name = get_chrom_name(in_file)
    if chrom_name is None:
        print("ERROR: Something went wrong reading your VCF file: a CHROM column could not be found. "
              "Please check the file is valid VCF format.")
        sys.exit(1)

    exclude = [chrom_name + "\t" + str(pos) for pos in mask_sites]
    temp_mask_file = in_file + "_maskTemp"
    with open_file(temp_mask_file, 'w') as fh:
        fh.write("\n".join(exclude))
    cleanup_files.append(temp_mask_file)

    #Read in/write out according to file ending
    in_call = "--gzvcf" if in_file.lower().endswith(".gz") else "--vcf"
    out_call = "| gzip -c" if out_file.lower().endswith(".gz") else ""

    call = ["vcftools", "--exclude-positions", shquote(temp_mask_file), in_call, shquote(in_file), "--recode --stdout", out_call, ">", shquote(out_file)]
    print("Removing masked sites from VCF file using vcftools... this may take some time. Call:")
    print(" ".join(call))
    run_shell_command(" ".join(call), raise_errors = True)
    # remove vcftools log file
    if cleanup:
        for file in cleanup_files:
            try:
                os.remove(file)
            except OSError:
                pass

def mask_fasta(mask_sites, in_file, out_file, mask_from_beginning=0, mask_from_end=0):
    """Mask the provided site list from a FASTA file and write to a new file.

    Masked sites are overwritten as "N"s.

:
    -----------
    mask_sites: list[int]
        A list of site indexes to exclude from the FASTA.
    in_file: str
        The path to the FASTA file you wish to mask.
    out_file: str
        The path to write the resulting FASTA to
    mask_from_beginning: int
       Number of sites to mask from the beginning of each sequence (default 0)
    mask_from_end: int
       Number of sites to mask from the end of each sequence (default 0)
    """
    # Load alignment as FASTA generator to prevent loading the whole alignment
    # into memory.
    alignment = SeqIO.parse(in_file, "fasta")

    # Write the masked alignment to disk one record at a time.
    print("Removing masked sites from FASTA file.")
    with open_file(out_file, "w") as oh:
        for record in alignment:
            # Convert to a mutable sequence to enable masking with Ns.
            sequence = MutableSeq(
                "N" * mask_from_beginning +
                str(record.seq)[mask_from_beginning:-mask_from_end or None] +
                "N" * mask_from_end
            )
            sequence_length = len(sequence)
            # Replace all excluded sites with Ns.
            for site in mask_sites:
                if site < sequence_length:
                    sequence[site] = "N"
            record.seq = sequence
            SeqIO.write(record, oh, "fasta")

def register_arguments(parser):
    parser.add_argument('--sequences', '-s', required=True, help="sequences in VCF or FASTA format")
    parser.add_argument('--mask', dest="mask_file", required=False, help="locations to be masked in BED file format")
    parser.add_argument('--output', '-o', help="output file")
    parser.add_argument('--no-cleanup', dest="cleanup", action="store_false",
                        help="Leave intermediate files around. May be useful for debugging")

def run(args):
    '''
    Mask specified sites from the VCF or FASTA.

    For VCF files, his occurs by removing them entirely from the VCF, essentially making
    them identical to the reference at the locations.

    For FASTA files, masked sites are replaced with "N".

    If users don't specify output, will overwrite the input file.
    '''
    # Check files exist and are not empty
    if not os.path.isfile(args.sequences):
        print("ERROR: File {} does not exist!".format(args.sequences))
        return 1
    if os.path.getsize(args.sequences) == 0:
        print("ERROR: {} is empty. Please check how this file was produced. "
              "Did an error occur in an earlier step?".format(args.sequences))
        return 1
    if args.mask_file:
        if not os.path.isfile(args.mask_file):
            print("ERROR: File {} does not exist!".format(args.mask_file))
        return 1
        if os.path.getsize(args.mask_file) == 0:
            print("ERROR: {} is an empty file.".format(args.mask_file))
            return 1

    mask_sites = []
    if args.mask_file:
        mask_sites.extend(read_bed_file(args.mask_file))

    # For both FASTA and VCF masking, we need a proper separate output file
    if args.output is not None:
        out_file = args.output
    else:
        out_file = os.path.join(os.path.dirname(args.sequences),
                                "masked_" + os.path.basename(args.sequences))

    if is_vcf(args.sequences):
        mask_vcf(mask_sites, args.sequences, out_file, args.cleanup)
    else:
        mask_fasta(mask_sites, args.sequences, out_file)

    if args.output is None:
        copyfile(out_file, args.sequences)
        if args.cleanup:
            os.remove(out_file)
<|endoftext|>


<commit_before>import asyncio, hashlib, logging, urllib

import urllib.request as urllib2
from http.cookiejar import CookieJar

from random import randrange, randint

import hangups

import plugins

logger = logging.getLogger(__name__)

__cleverbots = dict()

""" Cleverbot API adapted from https://github.com/folz/cleverbot.py """
class Cleverbot:
    """
    Wrapper over the Cleverbot API.
    """
    HOST = "www.cleverbot.com"
    PROTOCOL = "http://"
    RESOURCE = "/webservicemin"
    API_URL = PROTOCOL + HOST + RESOURCE

    headers = {
        'User-Agent': 'Mozilla/4.0 (compatible; MSIE 8.0; Windows NT 6.0)',
        'Accept': 'text/html,application/xhtml+xml,'
                  'application/xml;q=0.9,*/*;q=0.8',
        'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.7',
        'Accept-Language': 'en-us,en;q=0.8,en-us;q=0.5,en;q=0.3',
        'Cache-Control': 'no-cache',
        'Host': HOST,
        'Referer': PROTOCOL + HOST + '/',
        'Pragma': 'no-cache'
    }

    asked = 0
    lastanswer = ""

    def __init__(self):
        """ The data that will get passed to Cleverbot's web API """
        self.data = {
            'stimulus': '',
            'start': 'y',  # Never modified
            'sessionid': '',
            'vText8': '',
            'vText7': '',
            'vText6': '',
            'vText5': '',
            'vText4': '',
            'vText3': '',
            'vText2': '',
            'icognoid': 'wsf',  # Never modified
            'icognocheck': '',
            'fno': 0,  # Never modified
            'prevref': '',
            'emotionaloutput': '',  # Never modified
            'emotionalhistory': '',  # Never modified
            'asbotname': '',  # Never modified
            'ttsvoice': '',  # Never modified
            'typing': '',  # Never modified
            'lineref': '',
            'sub': 'Say',  # Never modified
            'islearning': 1,  # Never modified
            'cleanslate': False,  # Never modified
        }

        # the log of our conversation with Cleverbot
        self.conversation = []
        self.resp = str()
        self.asked = 0

        # install an opener with support for cookies
        cookies = CookieJar()
        handlers = [
            urllib2.HTTPHandler(),
            urllib2.HTTPSHandler(),
            urllib2.HTTPCookieProcessor(cookies)
        ]
        opener = urllib2.build_opener(*handlers)
        urllib2.install_opener(opener)

        # get the main page to get a cookie (see bug  #13)
        try:
            urllib2.urlopen(Cleverbot.PROTOCOL + Cleverbot.HOST)
        except urllib2.HTTPError:
            logger.exception(e)

    def ask(self, question):
        """Asks Cleverbot a question.

        Maintains message history.

        Args:
            q (str): The question to ask
        Returns:
            Cleverbot's answer
        """
        # Set the current question
        question = question.strip()
        if not question:
            return

        if not question.endswith(("!", ",", ".", ")", "%", "*")):
            # end a sentence with a full stop
            question += "."

        self.data['stimulus'] = question
        self.asked = self.asked + 1

        # Connect to Cleverbot's API and remember the response
        try:
            self.resp = self._send()
        except urllib2.HTTPError:
            # request failed. returning empty string
            return str()

        # Add the current question to the conversation log
        self.conversation.append(question)

        parsed = self._parse()

        # Set data as appropriate
        if not self.data['sessionid']:
            self.data['sessionid'] = parsed['conversation_id']

        # Add Cleverbot's reply to the conversation log
        self.conversation.append(parsed['answer'])
        self.lastanswer = parsed['answer']

        return parsed['answer']

    def _send(self):
        """POST the user's question and all required information to the
        Cleverbot API
        Cleverbot tries to prevent unauthorized access to its API by
        obfuscating how it generates the 'icognocheck' token, so we have
        to URLencode the data twice: once to generate the token, and
        twice to add the token to the data we're sending to Cleverbot.
        """
        # Set data as appropriate
        if self.conversation:
            linecount = 1
            for line in reversed(self.conversation):
                linecount += 1
                self.data['vText' + str(linecount)] = line
                if linecount == 8:
                    break

        """XXX: unlike the original code which used an unordered dict to
        build the payload, we use an ordered string to mimic the observed
        payload during normal operation of the bot
        """
        if self.asked <= 1:
            payload = ( "stimulus={0[stimulus]}"
                        "&cb_settings_scripting=no"
                        "&islearning=1"
                        "&icognoid={0[icognoid]}" ).format(self.data)

            query_string = ""

        else:
            payload = ( "stimulus={0[stimulus]}"
                        "&vText2={0[vText2]}"
                        "&vText3={0[vText3]}"
                        "&vText4={0[vText4]}"
                        "&vText5={0[vText5]}"
                        "&vText6={0[vText6]}"
                        "&vText7={0[vText7]}"
                        "&sessionid={0[sessionid]}"
                        "&cb_settings_language=es"
                        "&cb_settings_scripting=no"
                        "&islearning={0[islearning]}"
                        "&icognoid={0[icognoid]}" ).format(self.data)

            query_string = ("out={2}"
                            "&in={0[stimulus]}"
                            "&bot=c"
                            "&cbsid={0[sessionid]}"
                            "&xai={4}"
                            "&ns={1}"
                            "&al="
                            "&dl="
                            "&flag="
                            "&user="
                            "&mode=1"
                            "&t={3}"
                            "&").format(self.data,
                                        self.asked,
                                        self.lastanswer,
                                        randint(10000, 99999),
                                        self.data["sessionid"][0:3])

        # Generate the token
        digest_txt = payload[9:35]
        token = hashlib.md5(digest_txt.encode('utf-8')).hexdigest()
        payload += "&icognocheck={}".format(token)

        # Add the token to the data
        payload = payload.encode('utf-8')
        full_url = self.API_URL + "?" + query_string
        logger.debug(payload)
        logger.debug(full_url)
        req = urllib2.Request(full_url, payload, self.headers)

        # POST the data to Cleverbot's API
        conn = urllib2.urlopen(req)
        resp = conn.read()

        # Return Cleverbot's response
        return resp

    def _parse(self):
        """Parses Cleverbot's response"""
        parsed = [
            item.split('\r') for item in self.resp.decode('utf-8').split('\r\r\r\r\r\r')[:-1]
        ]
        parsed_dict = {
            'answer': parsed[0][0],
            'conversation_id': parsed[0][1],
            'conversation_log_id': parsed[0][2],
        }
        try:
            parsed_dict['unknown'] = parsed[1][-1]
        except IndexError:
            parsed_dict['unknown'] = None
        return parsed_dict

def _initialise(bot):
    plugins.register_handler(_handle_incoming_message, type="message")
    plugins.register_user_command(["chat"])
    plugins.register_admin_command(["chatreset"])

@asyncio.coroutine
def _handle_incoming_message(bot, event, command):
    """Handle random message intercepting"""

    if not event.text:
        return

    if not bot.get_config_suboption(event.conv_id, 'cleverbot_percentage_replies'):
        return

    percentage = bot.get_config_suboption(event.conv_id, 'cleverbot_percentage_replies')

    if randrange(0, 101, 1) < float(percentage):
        text = yield from cleverbot_ask(event.conv_id, event.text)
        if text:
            yield from bot.coro_send_message(event.conv_id, text)

def chat(bot, event, *args):
    """chat to Cleverbot"""
    if args:
        text = yield from cleverbot_ask(event.conv_id, ' '.join(args))
        if not text:
            text = _("Cleverbot is silent")

    else:
        text = _("you have to say something to Cleverbot")

    yield from bot.coro_send_message(event.conv_id, text)

def chatreset(bot, event, *args):
    if len(args) == 0:
        conv_id = event.conv_id
    else:
        conv_id = args[0]

    message = "no change"
    if conv_id in __cleverbots:
        del __cleverbots[conv_id]
        logger.debug("removed api instance for {}".format(conv_id))
        message = "removed {}".format(conv_id)

    yield from bot.coro_send_message(event.conv_id, message)

def cleverbot_ask(conv_id, message, filter_ads=True):
    if conv_id not in __cleverbots:
        __cleverbots[conv_id] = Cleverbot()
        logger.debug("added api instance for {}".format(conv_id))

    loop = asyncio.get_event_loop()

    text = False
    try:
        text = yield from loop.run_in_executor(None, __cleverbots[conv_id].ask, message)
        logger.debug("API returned: {}".format(text))
        if text:
            if filter_ads:
                if text.startswith("\n"):
                    # some ads appear to start with line breaks
                    text = False
                else:
                    # filter out specific ad-related keywords
                    ad_text = ["cleverscript", "cleverme", "clevertweet", "cleverenglish"]
                    for ad in ad_text:
                        if ad.lower() in text.lower():
                            logger.debug("ad-blocked")
                            text = False
                            break

    except:
        logger.exception("failed to get response")

    return text
<commit_msg>bugfix: accented characters for cleverbot replies<commit_after>import asyncio, hashlib, html, logging, urllib

import urllib.request as urllib2
from http.cookiejar import CookieJar

from random import randrange, randint

import hangups

import plugins

logger = logging.getLogger(__name__)

__cleverbots = dict()

""" Cleverbot API adapted from https://github.com/folz/cleverbot.py """
class Cleverbot:
    """
    Wrapper over the Cleverbot API.
    """
    HOST = "www.cleverbot.com"
    PROTOCOL = "http://"
    RESOURCE = "/webservicemin"
    API_URL = PROTOCOL + HOST + RESOURCE

    headers = {
        'User-Agent': 'Mozilla/4.0 (compatible; MSIE 8.0; Windows NT 6.0)',
        'Accept': 'text/html,application/xhtml+xml,'
                  'application/xml;q=0.9,*/*;q=0.8',
        'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.7',
        'Accept-Language': 'en-us,en;q=0.8,en-us;q=0.5,en;q=0.3',
        'Cache-Control': 'no-cache',
        'Host': HOST,
        'Referer': PROTOCOL + HOST + '/',
        'Pragma': 'no-cache'
    }

    asked = 0
    lastanswer = ""

    def __init__(self):
        """ The data that will get passed to Cleverbot's web API """
        self.data = {
            'stimulus': '',
            'start': 'y',  # Never modified
            'sessionid': '',
            'vText8': '',
            'vText7': '',
            'vText6': '',
            'vText5': '',
            'vText4': '',
            'vText3': '',
            'vText2': '',
            'icognoid': 'wsf',  # Never modified
            'icognocheck': '',
            'fno': 0,  # Never modified
            'prevref': '',
            'emotionaloutput': '',  # Never modified
            'emotionalhistory': '',  # Never modified
            'asbotname': '',  # Never modified
            'ttsvoice': '',  # Never modified
            'typing': '',  # Never modified
            'lineref': '',
            'sub': 'Say',  # Never modified
            'islearning': 1,  # Never modified
            'cleanslate': False,  # Never modified
        }

        # the log of our conversation with Cleverbot
        self.conversation = []
        self.resp = str()
        self.asked = 0

        # install an opener with support for cookies
        cookies = CookieJar()
        handlers = [
            urllib2.HTTPHandler(),
            urllib2.HTTPSHandler(),
            urllib2.HTTPCookieProcessor(cookies)
        ]
        opener = urllib2.build_opener(*handlers)
        urllib2.install_opener(opener)

        # get the main page to get a cookie (see bug  #13)
        try:
            urllib2.urlopen(Cleverbot.PROTOCOL + Cleverbot.HOST)
        except urllib2.HTTPError:
            logger.exception(e)

    def ask(self, question):
        """Asks Cleverbot a question.

        Maintains message history.

        Args:
            q (str): The question to ask
        Returns:
            Cleverbot's answer
        """
        # Set the current question
        question = question.strip()
        if not question:
            return

        if not question.endswith(("!", ",", ".", ")", "%", "*")):
            # end a sentence with a full stop
            question += "."

        question = question.encode("ascii", "xmlcharrefreplace")

        self.data['stimulus'] = question
        self.asked = self.asked + 1

        # Connect to Cleverbot's API and remember the response
        try:
            self.resp = self._send()
        except urllib2.HTTPError:
            # request failed. returning empty string
            return str()

        # Add the current question to the conversation log
        self.conversation.append(question)

        parsed = self._parse()

        # Set data as appropriate
        if not self.data['sessionid']:
            self.data['sessionid'] = parsed['conversation_id']

        # Add Cleverbot's reply to the conversation log
        self.conversation.append(parsed['answer'])
        self.lastanswer = parsed['answer']

        return parsed['answer']

    def _send(self):
        """POST the user's question and all required information to the
        Cleverbot API
        Cleverbot tries to prevent unauthorized access to its API by
        obfuscating how it generates the 'icognocheck' token, so we have
        to URLencode the data twice: once to generate the token, and
        twice to add the token to the data we're sending to Cleverbot.
        """
        # Set data as appropriate
        if self.conversation:
            linecount = 1
            for line in reversed(self.conversation):
                linecount += 1
                self.data['vText' + str(linecount)] = line
                if linecount == 8:
                    break

        """XXX: unlike the original code which used an unordered dict to
        build the payload, we use an ordered string to mimic the observed
        payload during normal operation of the bot
        """
        if self.asked <= 1:
            payload = ( "stimulus={0[stimulus]}"
                        "&cb_settings_scripting=no"
                        "&islearning=1"
                        "&icognoid={0[icognoid]}" ).format(self.data)

            query_string = ""

        else:
            payload = ( "stimulus={0[stimulus]}"
                        "&vText2={0[vText2]}"
                        "&vText3={0[vText3]}"
                        "&vText4={0[vText4]}"
                        "&vText5={0[vText5]}"
                        "&vText6={0[vText6]}"
                        "&vText7={0[vText7]}"
                        "&sessionid={0[sessionid]}"
                        "&cb_settings_language=es"
                        "&cb_settings_scripting=no"
                        "&islearning={0[islearning]}"
                        "&icognoid={0[icognoid]}" ).format(self.data)

            query_string = ("out={2}"
                            "&in={0[stimulus]}"
                            "&bot=c"
                            "&cbsid={0[sessionid]}"
                            "&xai={4}"
                            "&ns={1}"
                            "&al="
                            "&dl="
                            "&flag="
                            "&user="
                            "&mode=1"
                            "&t={3}"
                            "&").format(self.data,
                                        self.asked,
                                        self.lastanswer,
                                        randint(10000, 99999),
                                        self.data["sessionid"][0:3])

        # Generate the token
        digest_txt = payload[9:35]
        token = hashlib.md5(digest_txt.encode('utf-8')).hexdigest()
        payload += "&icognocheck={}".format(token)

        # Add the token to the data
        payload = payload.encode('utf-8')
        full_url = self.API_URL + "?" + query_string
        logger.debug(payload)
        logger.debug(full_url)
        req = urllib2.Request(full_url, payload, self.headers)

        # POST the data to Cleverbot's API
        conn = urllib2.urlopen(req)
        resp = conn.read()

        # Return Cleverbot's response
        return resp

    def _parse(self):
        """Parses Cleverbot's response"""
        parsed = [
            item.split('\r') for item in self.resp.decode('utf-8').split('\r\r\r\r\r\r')[:-1]
        ]
        parsed_dict = {
            'answer': parsed[0][0],
            'conversation_id': parsed[0][1],
            'conversation_log_id': parsed[0][2],
        }
        try:
            parsed_dict['unknown'] = parsed[1][-1]
        except IndexError:
            parsed_dict['unknown'] = None
        return parsed_dict

def _initialise(bot):
    plugins.register_handler(_handle_incoming_message, type="message")
    plugins.register_user_command(["chat"])
    plugins.register_admin_command(["chatreset"])

@asyncio.coroutine
def _handle_incoming_message(bot, event, command):
    """Handle random message intercepting"""

    if not event.text:
        return

    if not bot.get_config_suboption(event.conv_id, 'cleverbot_percentage_replies'):
        return

    percentage = bot.get_config_suboption(event.conv_id, 'cleverbot_percentage_replies')

    if randrange(0, 101, 1) < float(percentage):
        text = yield from cleverbot_ask(event.conv_id, event.text)
        if text:
            yield from bot.coro_send_message(event.conv_id, text)

def chat(bot, event, *args):
    """chat to Cleverbot"""
    if args:
        text = yield from cleverbot_ask(event.conv_id, ' '.join(args))
        if not text:
            text = _("Cleverbot is silent")

    else:
        text = _("you have to say something to Cleverbot")

    yield from bot.coro_send_message(event.conv_id, text)

def chatreset(bot, event, *args):
    if len(args) == 0:
        conv_id = event.conv_id
    else:
        conv_id = args[0]

    message = "no change"
    if conv_id in __cleverbots:
        del __cleverbots[conv_id]
        logger.debug("removed api instance for {}".format(conv_id))
        message = "removed {}".format(conv_id)

    yield from bot.coro_send_message(event.conv_id, message)

def cleverbot_ask(conv_id, message, filter_ads=True):
    if conv_id not in __cleverbots:
        __cleverbots[conv_id] = Cleverbot()
        logger.debug("added api instance for {}".format(conv_id))

    loop = asyncio.get_event_loop()

    text = False
    try:
        text = yield from loop.run_in_executor(None, __cleverbots[conv_id].ask, message)
        text = html.unescape(text)
        logger.debug("API returned: {}".format(text))
        if text:
            if filter_ads:
                if text.startswith("\n"):
                    # some ads appear to start with line breaks
                    text = False
                else:
                    # filter out specific ad-related keywords
                    ad_text = ["cleverscript", "cleverme", "clevertweet", "cleverenglish"]
                    for ad in ad_text:
                        if ad.lower() in text.lower():
                            logger.debug("ad-blocked")
                            text = False
                            break

    except:
        logger.exception("failed to get response")

    return text
<|endoftext|>


<commit_before>#!/usr/bin/env python

from cpsdirector import common, db

from conpaas.core.https import x509

import os
import platform

from distutils.spawn import find_executable

CERT_DIR = common.config.get('conpaas', 'CERT_DIR')
hostname = common.rlinput('Please enter your hostname: ', platform.node())

# create CA keypair
cakey = x509.gen_rsa_keypair()

# save ca_key.pem to filesystem
open(os.path.join(CERT_DIR, 'ca_key.pem'), 'w').write(x509.key_as_pem(cakey))

# create cert request
req = x509.create_x509_req(cakey, CN='CA', emailAddress='info@conpaas.eu',
    O='ConPaaS')

# create ca certificate, valid for five years
cacert = x509.create_cert(req, req, cakey, 1, 0, 60 * 60 * 24 * 365 * 5)

# save ca_cert.pem to filesystem
open(os.path.join(CERT_DIR, 'ca_cert.pem'), 'w').write(
    x509.cert_as_pem(cacert))

# create director key
dkey = x509.gen_rsa_keypair()

# save key.pem to filesystem
open(os.path.join(CERT_DIR, 'key.pem'), 'w').write(x509.key_as_pem(dkey))

# create director cert request
req = x509.create_x509_req(dkey, CN=hostname, emailAddress='info@conpaas.eu',
    O='ConPaaS', role='frontend')

# create director certificate
dcert = x509.create_cert(req, cacert, cakey, 1, 0, 60 * 60 * 24 * 365 * 5)

# save cert.pem to filesystem
open(os.path.join(CERT_DIR, 'cert.pem'), 'w').write(x509.cert_as_pem(dcert))

# find director.wsgi absolute path. We have to chdir or it will return the
# relative one
os.chdir('/')
wsgi_exec = find_executable('director.wsgi')

# create apache config file
conf_values = {
    'hostname': hostname,
    'wsgi':     wsgi_exec,
    'wsgidir':  os.path.dirname(wsgi_exec)
}

conf = """
<VirtualHost *>
    ServerName %(hostname)s

    WSGIDaemonProcess director user=www-data group=www-data threads=5
    WSGIScriptAlias / %(wsgi)s

    <Directory %(wsgidir)s>
        WSGIProcessGroup director
""" % conf_values

conf += """
        WSGIApplicationGroup %{GLOBAL}
        WSGIPassApacheRequest On

        SSLRequireSSL
        SSLVerifyClient optional
        SSLVerifyDepth 2
        SSLOptions +StdEnvVars +ExportCertData

        Order deny,allow
        Allow from all
    </Directory>

    SSLEngine on

    SSLCertificateFile    /etc/cpsdirector/certs/cert.pem
    SSLCertificateKeyFile /etc/cpsdirector/certs/key.pem

    SSLCACertificateFile  /etc/cpsdirector/certs/ca_cert.pem

    CustomLog ${APACHE_LOG_DIR}/director-access.log combined
    ErrorLog ${APACHE_LOG_DIR}/director-error.log
</VirtualHost>
"""

open('/etc/apache2/sites-available/conpaas-director', 'w').write(conf)

conffile = open(common.CONFFILE).read()
if 'DIRECTOR_URL' not in conffile:
    # append DIRECTOR_URL
    open(common.CONFFILE, 'a').write("\nDIRECTOR_URL = https://" + hostname)

db.create_all()
<commit_msg>conpaas-director: listen on port 5555<commit_after>#!/usr/bin/env python

from cpsdirector import common, db

from conpaas.core.https import x509

import os
import platform

from distutils.spawn import find_executable

CERT_DIR = common.config.get('conpaas', 'CERT_DIR')
hostname = common.rlinput('Please enter your hostname: ', platform.node())

# create CA keypair
cakey = x509.gen_rsa_keypair()

# save ca_key.pem to filesystem
open(os.path.join(CERT_DIR, 'ca_key.pem'), 'w').write(x509.key_as_pem(cakey))

# create cert request
req = x509.create_x509_req(cakey, CN='CA', emailAddress='info@conpaas.eu',
    O='ConPaaS')

# create ca certificate, valid for five years
cacert = x509.create_cert(req, req, cakey, 1, 0, 60 * 60 * 24 * 365 * 5)

# save ca_cert.pem to filesystem
open(os.path.join(CERT_DIR, 'ca_cert.pem'), 'w').write(
    x509.cert_as_pem(cacert))

# create director key
dkey = x509.gen_rsa_keypair()

# save key.pem to filesystem
open(os.path.join(CERT_DIR, 'key.pem'), 'w').write(x509.key_as_pem(dkey))

# create director cert request
req = x509.create_x509_req(dkey, CN=hostname, emailAddress='info@conpaas.eu',
    O='ConPaaS', role='frontend')

# create director certificate
dcert = x509.create_cert(req, cacert, cakey, 1, 0, 60 * 60 * 24 * 365 * 5)

# save cert.pem to filesystem
open(os.path.join(CERT_DIR, 'cert.pem'), 'w').write(x509.cert_as_pem(dcert))

# find director.wsgi absolute path. We have to chdir or it will return the
# relative one
os.chdir('/')
wsgi_exec = find_executable('director.wsgi')

# create apache config file
conf_values = {
    'hostname': hostname,
    'wsgi':     wsgi_exec,
    'wsgidir':  os.path.dirname(wsgi_exec),
    'port':     5555
}

conf = """
Listen %(port)s
<VirtualHost *:%(port)s>
    ServerName %(hostname)s

    WSGIDaemonProcess director user=www-data group=www-data threads=5
    WSGIScriptAlias / %(wsgi)s

    <Directory %(wsgidir)s>
        WSGIProcessGroup director
""" % conf_values

conf += """
        WSGIApplicationGroup %{GLOBAL}
        WSGIPassApacheRequest On

        SSLRequireSSL
        SSLVerifyClient optional
        SSLVerifyDepth 2
        SSLOptions +StdEnvVars +ExportCertData

        Order deny,allow
        Allow from all
    </Directory>

    SSLEngine on

    SSLCertificateFile    /etc/cpsdirector/certs/cert.pem
    SSLCertificateKeyFile /etc/cpsdirector/certs/key.pem

    SSLCACertificateFile  /etc/cpsdirector/certs/ca_cert.pem

    CustomLog ${APACHE_LOG_DIR}/director-access.log combined
    ErrorLog ${APACHE_LOG_DIR}/director-error.log
</VirtualHost>
"""

open('/etc/apache2/sites-available/conpaas-director', 'w').write(conf)

conffile = open(common.CONFFILE).read()
if 'DIRECTOR_URL' not in conffile:
    # append DIRECTOR_URL
    open(common.CONFFILE, 'a').write("\nDIRECTOR_URL = https://%s:%s" % (hostname, conf_values['port']))

db.create_all()
<|endoftext|>


<commit_before>from copy import copy
import ctypes
import threading
import traceback
import time

from datetime import datetime

from . import screencapturer, config
from .enumeration import TestFixtureStatus, TestClassRunMode, TestCaseStatus
from .plogger import preporter, pconsole
from .testsuite import AfterSuite, BeforeSuite, AfterClass, BeforeClass, BeforeGroup, AfterGroup, AfterMethod, \
    BeforeMethod
from .plistener import test_listeners

__author__ = 'karl.gong'

class TestExecutor(threading.Thread):
    def __init__(self, parent_test_executor):
        threading.Thread.__init__(self)
        self.parent_test_executor = parent_test_executor
        if self.parent_test_executor:
            self.__properties = copy(self.parent_test_executor.get_properties())
        else:
            self.__properties = {}

    def start_and_join(self):
        self.start()
        self.join()

    def update_properties(self, properties):
        self.__properties.update(properties)

    def clear_properties(self):
        self.__properties.clear()

    def get_property(self, key):
        try:
            return self.__properties[key]
        except KeyError:
            return None

    def get_properties(self):
        return self.__properties

class TestSuiteExecutor(TestExecutor):
    def __init__(self, test_suite):
        TestExecutor.__init__(self, None)
        self.test_suite = test_suite

    def run(self):
        test_listeners.on_test_suite_start(self.test_suite)
        self.test_suite.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_suite.before_suite).start_and_join()

        test_class_executors = []

        for test_class in self.test_suite.test_classes:
            test_class_executor = TestClassExecutor(self, test_class)
            test_class_executors.append(test_class_executor)
            test_class_executor.start()

        for executor in test_class_executors:
            executor.join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_suite.after_suite).start_and_join()
        self.test_suite.end_time = datetime.now()
        test_listeners.on_test_suite_finish(self.test_suite)

class TestClassExecutor(TestExecutor):
    def __init__(self, test_suite_executor, test_class):
        TestExecutor.__init__(self, test_suite_executor)
        self.test_class = test_class

    def run(self):
        test_listeners.on_test_class_start(self.test_class)
        self.test_class.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_class.before_class).start_and_join()

        if self.test_class.run_mode == TestClassRunMode.SingleLine:
            for test_group in self.test_class.test_groups:
                TestGroupExecutor(self, test_group).start_and_join()
        else:
            test_group_executors = []

            for test_group in self.test_class.test_groups:
                test_group_executor = TestGroupExecutor(self, test_group)
                test_group_executors.append(test_group_executor)
                test_group_executor.start()

            for executor in test_group_executors:
                executor.join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_class.after_class).start_and_join()
        self.test_class.end_time = datetime.now()
        test_listeners.on_test_class_finish(self.test_class)

class TestGroupExecutor(TestExecutor):
    def __init__(self, test_class_executor, test_group):
        TestExecutor.__init__(self, test_class_executor)
        self.test_group = test_group

    def run(self):
        test_listeners.on_test_group_start(self.test_group)
        self.test_group.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_group.before_group).start_and_join()

        if self.test_group.test_class.run_mode == TestClassRunMode.SingleLine:
            for test_case in self.test_group.test_cases:
                TestCaseExecutor(self, test_case).start_and_join()
        else:
            test_case_executors = []

            for test_case in self.test_group.test_cases:
                test_case_executor = TestCaseExecutor(self, test_case)
                test_case_executors.append(test_case_executor)
                test_case_executor.start()

            for executor in test_case_executors:
                executor.join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_group.after_group).start_and_join()
        self.test_group.end_time = datetime.now()
        test_listeners.on_test_group_finish(self.test_group)

class TestCaseExecutor(TestExecutor):
    def __init__(self, test_group_executor, test_case):
        TestExecutor.__init__(self, test_group_executor)
        self.test_case = test_case

    def run(self):
        test_listeners.on_test_case_start(self.test_case)
        self.test_case.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_case.before_method).start_and_join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_case.test).start_and_join()

        logger_filler = "-" * (100 - len(self.test_case.full_name) - 6)
        if self.test_case.status == TestCaseStatus.PASSED:
            pconsole.write_line("%s%s|PASS|" % (self.test_case.full_name, logger_filler))
        elif self.test_case.status == TestCaseStatus.FAILED:
            pconsole.write_line("%s%s|FAIL|" % (self.test_case.full_name, logger_filler))
        elif self.test_case.status == TestCaseStatus.SKIPPED:
            pconsole.write_line("%s%s|SKIP|" % (self.test_case.full_name, logger_filler))

        test_fixture_executor_pool.apply_for_executor(self, self.test_case.after_method).start_and_join()
        self.test_case.end_time = datetime.now()
        test_listeners.on_test_case_finish(self.test_case)

class TestFixtureExecutor(TestExecutor):
    def __init__(self, test_case_executor, test_fixture):
        TestExecutor.__init__(self, test_case_executor)
        self.test_fixture = test_fixture

    def run_test_fixture(self):
        if self.test_fixture.is_empty:
            self.test_fixture.status = TestFixtureStatus.PASSED
            return

        self.test_fixture.status = TestFixtureStatus.RUNNING
        test_fixture_sub_executor = TestFixtureSubExecutor(self)
        test_fixture_sub_executor.start()
        if self.test_fixture.timeout > 0:
            test_fixture_sub_executor.join(self.test_fixture.timeout)
            if test_fixture_sub_executor.isAlive():
                from .plogger import preporter
                from . import screencapturer
                self.test_fixture.status = TestFixtureStatus.FAILED
                self.test_fixture.failure_message = "Timed out executing this test fixture in %s seconds." % self.test_fixture.timeout
                self.test_fixture.failure_type = "TimeoutException"
                self.test_fixture.stack_trace = ""
                preporter.error("Failed with following message:\n" + self.test_fixture.failure_message)
                screencapturer.take_screenshot()
                kill_thread(test_fixture_sub_executor)
        else:
            test_fixture_sub_executor.join()

    def skip_test_fixture(self, caused_test_fixture):
        if self.test_fixture.is_empty:
            self.test_fixture.status = TestFixtureStatus.SKIPPED
            return

        from .plogger import preporter
        self.test_fixture.status = TestFixtureStatus.SKIPPED
        self.test_fixture.skip_message = "@%s failed, so skipped." % caused_test_fixture.fixture_type
        preporter.warn("@%s failed, so skipped." % caused_test_fixture.fixture_type)

    def run(self):
        self.test_fixture.start_time = datetime.now()
        self.update_properties({"running_test_fixture": self.test_fixture})

        failed_setup_fixture = self.test_fixture.context.get_failed_setup_fixture()
        if not failed_setup_fixture:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterSuite) and isinstance(failed_setup_fixture, BeforeSuite) and self.test_fixture.always_run:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterClass) and isinstance(failed_setup_fixture, BeforeClass) and self.test_fixture.always_run:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterGroup) and isinstance(failed_setup_fixture, BeforeGroup) and self.test_fixture.always_run:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterMethod) and isinstance(failed_setup_fixture, BeforeMethod) and self.test_fixture.always_run:
            self.run_test_fixture()
        else:
            self.skip_test_fixture(failed_setup_fixture)

        # spread before's attributes
        if isinstance(self.test_fixture, BeforeSuite):
            before_suite_dict = self.test_fixture.test_fixture_ref.__self__.__dict__
            for test_class in self.test_fixture.test_suite.test_classes:
                test_class.test_class_ref.__dict__.update(before_suite_dict)
                for test_group in test_class.test_groups:
                    test_group.test_class_ref.__dict__.update(before_suite_dict)
                    for test_case in test_group.test_cases:
                        test_case.test_case_ref.__self__.__dict__.update(before_suite_dict)
        elif isinstance(self.test_fixture, BeforeClass):
            before_class_dict = self.test_fixture.test_fixture_ref.__self__.__dict__
            for test_group in self.test_fixture.test_class.test_groups:
                test_group.test_class_ref.__dict__.update(before_class_dict)
                for test_case in test_group.test_cases:
                    test_case.test_case_ref.__self__.__dict__.update(before_class_dict)
        elif isinstance(self.test_fixture, BeforeGroup):
            before_group_dict = self.test_fixture.test_fixture_ref.__self__.__dict__
            for test_case in self.test_fixture.test_group.test_cases:
                test_case.test_case_ref.__self__.__dict__.update(before_group_dict)

        self.update_properties({"running_test_fixture": None})
        self.test_fixture.end_time = datetime.now()

class TestFixtureSubExecutor(TestExecutor):
    def __init__(self, parent_executor):
        TestExecutor.__init__(self, parent_executor)
        self.test_fixture = parent_executor.test_fixture
        self.setDaemon(True)

    def run(self):
        try:
            if self.test_fixture.arguments_count == 1:
                self.test_fixture.test_fixture_ref.__call__()
            elif self.test_fixture.arguments_count == 2:
                self.test_fixture.test_fixture_ref.__call__(self.test_fixture.context)
        except Exception as e:
            self.test_fixture.status = TestFixtureStatus.FAILED
            self.test_fixture.failure_message = "\n".join([str(arg) for arg in e.args])
            self.test_fixture.failure_type = e.__class__.__name__
            self.test_fixture.stack_trace = traceback.format_exc()
            preporter.error("Failed with following message:\n%s" % self.test_fixture.stack_trace)
            screencapturer.take_screenshot()
        else:
            self.test_fixture.status = TestFixtureStatus.PASSED

class TestFixtureExecutorPool(threading.Thread):
    def __init__(self, size):
        threading.Thread.__init__(self)
        self.size = size
        self.running_test_executors = []
        self.pending_test_fixtures = []
        self.__test_fixture_lock = threading.Lock()
        self.__test_executor_lock = threading.Lock()

    def get_top_test_fixture(self):
        self.__test_fixture_lock.acquire()
        try:
            self.pending_test_fixtures = sorted(self.pending_test_fixtures, key=lambda item: item.execution_priority)
            return self.pending_test_fixtures[0]
        finally:
            self.__test_fixture_lock.release()

    def get_running_test_executors_count(self):
        self.__test_executor_lock.acquire()
        try:
            self.running_test_executors = [test_executor for test_executor in self.running_test_executors if
                                           test_executor.isAlive()]
            return len(self.running_test_executors)
        finally:
            self.__test_executor_lock.release()

    def apply_for_executor(self, test_executor, test_fixture):
        self.__test_fixture_lock.acquire()
        try:
            self.pending_test_fixtures.append(test_fixture)
        finally:
            self.__test_fixture_lock.release()

        while self.get_running_test_executors_count() >= self.size or self.get_top_test_fixture() != test_fixture:
            time.sleep(1)

        test_fixture_executor = TestFixtureExecutor(test_executor, test_fixture)

        self.__test_executor_lock.acquire()
        try:
            self.running_test_executors.append(test_fixture_executor)
        finally:
            self.__test_executor_lock.release()

        self.__test_fixture_lock.acquire()
        try:
            self.pending_test_fixtures.remove(test_fixture)
        finally:
            self.__test_fixture_lock.release()

        return test_fixture_executor

test_fixture_executor_pool = TestFixtureExecutorPool(int(config.get_option("test_executor_number")))

def current_executor():
    return threading.currentThread()

def kill_thread(thread):
    """Kill a python thread from another thread.

    :param thread: a threading.Thread instance
    """
    exc = ctypes.py_object(SystemExit)
    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
        ctypes.c_long(thread.ident), exc)
    if res == 0:
        raise ValueError("nonexistent thread id")
    elif res > 1:
        # """if it returns a number greater than one, you're in trouble,
        # and you should call it again with exc=NULL to revert the effect"""
        ctypes.pythonapi.PyThreadState_SetAsyncExc(thread.ident, None)
        raise SystemError("PyThreadState_SetAsyncExc failed")
<commit_msg>Do not spread while test_fixture is empty.<commit_after>from copy import copy
import ctypes
import threading
import traceback
import time

from datetime import datetime

from . import screencapturer, config
from .enumeration import TestFixtureStatus, TestClassRunMode, TestCaseStatus
from .plogger import preporter, pconsole
from .testsuite import AfterSuite, BeforeSuite, AfterClass, BeforeClass, BeforeGroup, AfterGroup, AfterMethod, \
    BeforeMethod
from .plistener import test_listeners

__author__ = 'karl.gong'

class TestExecutor(threading.Thread):
    def __init__(self, parent_test_executor):
        threading.Thread.__init__(self)
        self.parent_test_executor = parent_test_executor
        if self.parent_test_executor:
            self.__properties = copy(self.parent_test_executor.get_properties())
        else:
            self.__properties = {}

    def start_and_join(self):
        self.start()
        self.join()

    def update_properties(self, properties):
        self.__properties.update(properties)

    def clear_properties(self):
        self.__properties.clear()

    def get_property(self, key):
        try:
            return self.__properties[key]
        except KeyError:
            return None

    def get_properties(self):
        return self.__properties

class TestSuiteExecutor(TestExecutor):
    def __init__(self, test_suite):
        TestExecutor.__init__(self, None)
        self.test_suite = test_suite

    def run(self):
        test_listeners.on_test_suite_start(self.test_suite)
        self.test_suite.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_suite.before_suite).start_and_join()

        test_class_executors = []

        for test_class in self.test_suite.test_classes:
            test_class_executor = TestClassExecutor(self, test_class)
            test_class_executors.append(test_class_executor)
            test_class_executor.start()

        for executor in test_class_executors:
            executor.join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_suite.after_suite).start_and_join()
        self.test_suite.end_time = datetime.now()
        test_listeners.on_test_suite_finish(self.test_suite)

class TestClassExecutor(TestExecutor):
    def __init__(self, test_suite_executor, test_class):
        TestExecutor.__init__(self, test_suite_executor)
        self.test_class = test_class

    def run(self):
        test_listeners.on_test_class_start(self.test_class)
        self.test_class.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_class.before_class).start_and_join()

        if self.test_class.run_mode == TestClassRunMode.SingleLine:
            for test_group in self.test_class.test_groups:
                TestGroupExecutor(self, test_group).start_and_join()
        else:
            test_group_executors = []

            for test_group in self.test_class.test_groups:
                test_group_executor = TestGroupExecutor(self, test_group)
                test_group_executors.append(test_group_executor)
                test_group_executor.start()

            for executor in test_group_executors:
                executor.join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_class.after_class).start_and_join()
        self.test_class.end_time = datetime.now()
        test_listeners.on_test_class_finish(self.test_class)

class TestGroupExecutor(TestExecutor):
    def __init__(self, test_class_executor, test_group):
        TestExecutor.__init__(self, test_class_executor)
        self.test_group = test_group

    def run(self):
        test_listeners.on_test_group_start(self.test_group)
        self.test_group.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_group.before_group).start_and_join()

        if self.test_group.test_class.run_mode == TestClassRunMode.SingleLine:
            for test_case in self.test_group.test_cases:
                TestCaseExecutor(self, test_case).start_and_join()
        else:
            test_case_executors = []

            for test_case in self.test_group.test_cases:
                test_case_executor = TestCaseExecutor(self, test_case)
                test_case_executors.append(test_case_executor)
                test_case_executor.start()

            for executor in test_case_executors:
                executor.join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_group.after_group).start_and_join()
        self.test_group.end_time = datetime.now()
        test_listeners.on_test_group_finish(self.test_group)

class TestCaseExecutor(TestExecutor):
    def __init__(self, test_group_executor, test_case):
        TestExecutor.__init__(self, test_group_executor)
        self.test_case = test_case

    def run(self):
        test_listeners.on_test_case_start(self.test_case)
        self.test_case.start_time = datetime.now()
        test_fixture_executor_pool.apply_for_executor(self, self.test_case.before_method).start_and_join()

        test_fixture_executor_pool.apply_for_executor(self, self.test_case.test).start_and_join()

        logger_filler = "-" * (100 - len(self.test_case.full_name) - 6)
        if self.test_case.status == TestCaseStatus.PASSED:
            pconsole.write_line("%s%s|PASS|" % (self.test_case.full_name, logger_filler))
        elif self.test_case.status == TestCaseStatus.FAILED:
            pconsole.write_line("%s%s|FAIL|" % (self.test_case.full_name, logger_filler))
        elif self.test_case.status == TestCaseStatus.SKIPPED:
            pconsole.write_line("%s%s|SKIP|" % (self.test_case.full_name, logger_filler))

        test_fixture_executor_pool.apply_for_executor(self, self.test_case.after_method).start_and_join()
        self.test_case.end_time = datetime.now()
        test_listeners.on_test_case_finish(self.test_case)

class TestFixtureExecutor(TestExecutor):
    def __init__(self, test_case_executor, test_fixture):
        TestExecutor.__init__(self, test_case_executor)
        self.test_fixture = test_fixture

    def run_test_fixture(self):
        if self.test_fixture.is_empty:
            self.test_fixture.status = TestFixtureStatus.PASSED
            return

        self.test_fixture.status = TestFixtureStatus.RUNNING
        test_fixture_sub_executor = TestFixtureSubExecutor(self)
        test_fixture_sub_executor.start()
        if self.test_fixture.timeout > 0:
            test_fixture_sub_executor.join(self.test_fixture.timeout)
            if test_fixture_sub_executor.isAlive():
                from .plogger import preporter
                from . import screencapturer
                self.test_fixture.status = TestFixtureStatus.FAILED
                self.test_fixture.failure_message = "Timed out executing this test fixture in %s seconds." % self.test_fixture.timeout
                self.test_fixture.failure_type = "TimeoutException"
                self.test_fixture.stack_trace = ""
                preporter.error("Failed with following message:\n" + self.test_fixture.failure_message)
                screencapturer.take_screenshot()
                kill_thread(test_fixture_sub_executor)
        else:
            test_fixture_sub_executor.join()

    def skip_test_fixture(self, caused_test_fixture):
        if self.test_fixture.is_empty:
            self.test_fixture.status = TestFixtureStatus.SKIPPED
            return

        from .plogger import preporter
        self.test_fixture.status = TestFixtureStatus.SKIPPED
        self.test_fixture.skip_message = "@%s failed, so skipped." % caused_test_fixture.fixture_type
        preporter.warn("@%s failed, so skipped." % caused_test_fixture.fixture_type)

    def run(self):
        self.test_fixture.start_time = datetime.now()
        self.update_properties({"running_test_fixture": self.test_fixture})

        failed_setup_fixture = self.test_fixture.context.get_failed_setup_fixture()
        if not failed_setup_fixture:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterSuite) and isinstance(failed_setup_fixture, BeforeSuite) and self.test_fixture.always_run:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterClass) and isinstance(failed_setup_fixture, BeforeClass) and self.test_fixture.always_run:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterGroup) and isinstance(failed_setup_fixture, BeforeGroup) and self.test_fixture.always_run:
            self.run_test_fixture()
        elif isinstance(self.test_fixture, AfterMethod) and isinstance(failed_setup_fixture, BeforeMethod) and self.test_fixture.always_run:
            self.run_test_fixture()
        else:
            self.skip_test_fixture(failed_setup_fixture)

        # spread before's attributes
        if not self.test_fixture.is_empty:
            if isinstance(self.test_fixture, BeforeSuite):
                before_suite_dict = self.test_fixture.test_fixture_ref.__self__.__dict__
                for test_class in self.test_fixture.test_suite.test_classes:
                    test_class.test_class_ref.__dict__.update(before_suite_dict)
                    for test_group in test_class.test_groups:
                        test_group.test_class_ref.__dict__.update(before_suite_dict)
                        for test_case in test_group.test_cases:
                            test_case.test_case_ref.__self__.__dict__.update(before_suite_dict)
            elif isinstance(self.test_fixture, BeforeClass):
                before_class_dict = self.test_fixture.test_fixture_ref.__self__.__dict__
                for test_group in self.test_fixture.test_class.test_groups:
                    test_group.test_class_ref.__dict__.update(before_class_dict)
                    for test_case in test_group.test_cases:
                        test_case.test_case_ref.__self__.__dict__.update(before_class_dict)
            elif isinstance(self.test_fixture, BeforeGroup):
                before_group_dict = self.test_fixture.test_fixture_ref.__self__.__dict__
                for test_case in self.test_fixture.test_group.test_cases:
                    test_case.test_case_ref.__self__.__dict__.update(before_group_dict)

        self.update_properties({"running_test_fixture": None})
        self.test_fixture.end_time = datetime.now()

class TestFixtureSubExecutor(TestExecutor):
    def __init__(self, parent_executor):
        TestExecutor.__init__(self, parent_executor)
        self.test_fixture = parent_executor.test_fixture
        self.setDaemon(True)

    def run(self):
        try:
            if self.test_fixture.arguments_count == 1:
                self.test_fixture.test_fixture_ref.__call__()
            elif self.test_fixture.arguments_count == 2:
                self.test_fixture.test_fixture_ref.__call__(self.test_fixture.context)
        except Exception as e:
            self.test_fixture.status = TestFixtureStatus.FAILED
            self.test_fixture.failure_message = "\n".join([str(arg) for arg in e.args])
            self.test_fixture.failure_type = e.__class__.__name__
            self.test_fixture.stack_trace = traceback.format_exc()
            preporter.error("Failed with following message:\n%s" % self.test_fixture.stack_trace)
            screencapturer.take_screenshot()
        else:
            self.test_fixture.status = TestFixtureStatus.PASSED

class TestFixtureExecutorPool(threading.Thread):
    def __init__(self, size):
        threading.Thread.__init__(self)
        self.size = size
        self.running_test_executors = []
        self.pending_test_fixtures = []
        self.__test_fixture_lock = threading.Lock()
        self.__test_executor_lock = threading.Lock()

    def get_top_test_fixture(self):
        self.__test_fixture_lock.acquire()
        try:
            self.pending_test_fixtures = sorted(self.pending_test_fixtures, key=lambda item: item.execution_priority)
            return self.pending_test_fixtures[0]
        finally:
            self.__test_fixture_lock.release()

    def get_running_test_executors_count(self):
        self.__test_executor_lock.acquire()
        try:
            self.running_test_executors = [test_executor for test_executor in self.running_test_executors if
                                           test_executor.isAlive()]
            return len(self.running_test_executors)
        finally:
            self.__test_executor_lock.release()

    def apply_for_executor(self, test_executor, test_fixture):
        self.__test_fixture_lock.acquire()
        try:
            self.pending_test_fixtures.append(test_fixture)
        finally:
            self.__test_fixture_lock.release()

        while self.get_running_test_executors_count() >= self.size or self.get_top_test_fixture() != test_fixture:
            time.sleep(1)

        test_fixture_executor = TestFixtureExecutor(test_executor, test_fixture)

        self.__test_executor_lock.acquire()
        try:
            self.running_test_executors.append(test_fixture_executor)
        finally:
            self.__test_executor_lock.release()

        self.__test_fixture_lock.acquire()
        try:
            self.pending_test_fixtures.remove(test_fixture)
        finally:
            self.__test_fixture_lock.release()

        return test_fixture_executor

test_fixture_executor_pool = TestFixtureExecutorPool(int(config.get_option("test_executor_number")))

def current_executor():
    return threading.currentThread()

def kill_thread(thread):
    """Kill a python thread from another thread.

    :param thread: a threading.Thread instance
    """
    exc = ctypes.py_object(SystemExit)
    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
        ctypes.c_long(thread.ident), exc)
    if res == 0:
        raise ValueError("nonexistent thread id")
    elif res > 1:
        # """if it returns a number greater than one, you're in trouble,
        # and you should call it again with exc=NULL to revert the effect"""
        ctypes.pythonapi.PyThreadState_SetAsyncExc(thread.ident, None)
        raise SystemError("PyThreadState_SetAsyncExc failed")
<|endoftext|>


<commit_before>from django.db import models

from datetime import date, datetime, timedelta
from mechanize import Browser
import re

from grid.scoreboard_parser import ScoreboardParser

class Scraper(models.Model):
    run_at = models.DateTimeField()

    def scrape(self):
        if self.run_at > datetime.now():
            return

        week = Week.objects.latest()
        all_done = True
        for game in week.game_set.all():
            if not game.is_final():
                all_done = False
                break
        if all_done:
            return

        br = Browser()
        res = br.open(week.scoreboard_url)
        content = res.read()
        parser = ScoreboardParser()
        for table in re.findall('<table class="data.*?
', content):
            parser.feed(table)
        next_run = datetime.now() + timedelta(days=1)

        for parsed_game in parser.scores:
            if len(parsed_game) == 5:
                time_left = parsed_game[0].strip()
                away_team = parsed_game[1].strip()
                away_score = parsed_game[2].strip()
                home_team = parsed_game[3].strip()
                home_score = parsed_game[4].strip()
            else:
                time_left = parsed_game[0].strip()
                away_team = parsed_game[1].strip()
                away_score = 0
                home_team = parsed_game[2].strip()
                home_score = 0

            game = week.game_set.get(away_team=away_team)
            game.away_score = away_score
            game.home_score = home_score
            game.time_left = time_left
            game.save()

            if game.in_progress():
                next_run = datetime.now()
            elif not game.is_final():
                now = datetime.now()
                kickoff = datetime.strptime(game.time_left, "%I:%M %p")
                gametime = datetime(year=now.year, month=now.month, day=now.day, hour=kickoff.hour-1, minute=kickoff.minute)

                next_run = min(next_run, gametime)
        self.run_at = next_run
        self.save()

class Week(models.Model):
    start_date = models.DateField('Date')
    number = models.IntegerField()
    scoreboard_url = models.CharField(max_length=255)

    def __unicode__(self):
        return "Week %d" % self.number

    class Meta:
        get_latest_by = 'start_date'

class Game(models.Model):
    PICKER_CHOICES = (
        ("BEN", "Ben"),
        ("BRIAN", "Brian"),
    )
    home_team = models.CharField(max_length=3)
    home_score = models.IntegerField()
    away_team = models.CharField(max_length=3)
    away_score = models.IntegerField()

    week = models.ForeignKey(Week)

    picked_team = models.CharField(max_length=3)
    picker = models.CharField(max_length=5, choices=PICKER_CHOICES)
    lock = models.BooleanField(default=False)
    anti_lock = models.BooleanField(default=False)

    time_left = models.CharField(max_length=100)

    def __unicode__(self):
        return "%s @ %s - Week %d" % (self.away_team, self.home_team, self.week.number)

    def is_final(self):
        return self.time_left == "Final"

    def in_progress(self):
        if "QTR" in self.time_left:
            return True
        elif "OT" in self.time_left:
            return True
        elif "Halftime" in self.time_left:
            return True
        return False
<commit_msg>Handle games that end in overtime.<commit_after>from django.db import models

from datetime import date, datetime, timedelta
from mechanize import Browser
import re

from grid.scoreboard_parser import ScoreboardParser

class Scraper(models.Model):
    run_at = models.DateTimeField()

    def scrape(self):
        if self.run_at > datetime.now():
            return

        week = Week.objects.latest()
        all_done = True
        for game in week.game_set.all():
            if not game.is_final():
                all_done = False
                break
        if all_done:
            return

        br = Browser()
        res = br.open(week.scoreboard_url)
        content = res.read()
        parser = ScoreboardParser()
        for table in re.findall('<table class="data.*?
', content):
            parser.feed(table)
        next_run = datetime.now() + timedelta(days=1)

        for parsed_game in parser.scores:
            if len(parsed_game) == 5:
                time_left = parsed_game[0].strip()
                away_team = parsed_game[1].strip()
                away_score = parsed_game[2].strip()
                home_team = parsed_game[3].strip()
                home_score = parsed_game[4].strip()
            else:
                time_left = parsed_game[0].strip()
                away_team = parsed_game[1].strip()
                away_score = 0
                home_team = parsed_game[2].strip()
                home_score = 0

            if time_left == "Final OT":
                time_left = "Final"
            game = week.game_set.get(away_team=away_team)
            game.away_score = away_score
            game.home_score = home_score
            game.time_left = time_left
            game.save()

            if game.in_progress():
                next_run = datetime.now()
            elif not game.is_final():
                now = datetime.now()
                kickoff = datetime.strptime(game.time_left, "%I:%M %p")
                gametime = datetime(year=now.year, month=now.month, day=now.day, hour=kickoff.hour-1, minute=kickoff.minute)

                next_run = min(next_run, gametime)
        self.run_at = next_run
        self.save()

class Week(models.Model):
    start_date = models.DateField('Date')
    number = models.IntegerField()
    scoreboard_url = models.CharField(max_length=255)

    def __unicode__(self):
        return "Week %d" % self.number

    class Meta:
        get_latest_by = 'start_date'

class Game(models.Model):
    PICKER_CHOICES = (
        ("BEN", "Ben"),
        ("BRIAN", "Brian"),
    )
    home_team = models.CharField(max_length=3)
    home_score = models.IntegerField()
    away_team = models.CharField(max_length=3)
    away_score = models.IntegerField()

    week = models.ForeignKey(Week)

    picked_team = models.CharField(max_length=3)
    picker = models.CharField(max_length=5, choices=PICKER_CHOICES)
    lock = models.BooleanField(default=False)
    anti_lock = models.BooleanField(default=False)

    time_left = models.CharField(max_length=100)

    def __unicode__(self):
        return "%s @ %s - Week %d" % (self.away_team, self.home_team, self.week.number)

    def is_final(self):
        return self.time_left == "Final"

    def in_progress(self):
        if "QTR" in self.time_left:
            return True
        elif "OT" in self.time_left:
            return True
        elif "Halftime" in self.time_left:
            return True
        return False
<|endoftext|>


<commit_before>#!/usr/bin/python

import RPi.GPIO as GPIO
import nrf24
import sys
import time
import spidev
import argparse
import os
import sqlite3
import datetime
import errno
import socket

PIPES = ([0xe7, 0xe7, 0xe7, 0xe7, 0xe7], [0xc2, 0xc2, 0xc2, 0xc2, 0xc2])

class Relay:
    def __init__(self, pins):
        self.pins = pins
        for pin in self.pins.values():
            GPIO.setup(pin, GPIO.OUT, initial=GPIO.HIGH)

    def output(self, pin, state):
        print "setting pin", pin, state and "on" or "off"
        GPIO.output(self.pins[pin], not state)  # These devices are active-low.

    def cleanup(self):
        pass  # this will be done later: GPIO.cleanup()

class Temperature:
    def __init__(self, major=0, minor=0):
        self.spi = spidev.SpiDev()
        self.spi.open(major, minor)

    def rawread(self):
        return self.spi.xfer2([0, 0])

    def read(self):
        return self.calc_temp(self.rawread())

    @staticmethod
    def calc_temp(buf):
        return (((buf[0] << 8) | buf[1]) >> 3) * 0.0625

    def cleanup(self):
        self.spi.close()

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.cleanup()

class Boiler:
    def __init__(self, major, minor, ce_pin, irq_pin, temperature, relay):
        self.relay = relay
        self.temperature = temperature
        self.radio = nrf24.NRF24()
        self.radio.begin(major, minor, ce_pin, irq_pin)
        self.radio.setDataRate(self.radio.BR_250KBPS)
        self.radio.setChannel(0)
        self.radio.setAutoAck(1)
        self.radio.enableDynamicPayloads()
        self.radio.printDetails()
        self.radio.openWritingPipe(PIPES[0])
        self.radio.openReadingPipe(1, PIPES[1])

    def run(self):
        while True:
            try:
                self.radio.startListening()
                recv_buffer = self.recv(10)
                self.radio.stopListening()
                for byte in recv_buffer:
                    pin = byte >> 1
                    state = byte & 0x1
                    self.relay.output(pin, state)
                result = self.radio.write(self.temperature.rawread())
                if not result:
                    print datetime.datetime.now(), "Did not receive ACK from controller."
            except Exception as e:
                print e

    def recv(self, timeout=None):
        end = time.time() + timeout
        pipe = [0]
        while not self.radio.available(pipe) and (timeout is None or time.time() < end):
            time.sleep(10000 / 1e6)
        if self.radio.available(pipe):
            recv_buffer = []
            self.radio.read(recv_buffer)
            return recv_buffer
        return []

    def cleanup(self):
        self.radio.end()

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.cleanup()

class Controller:
    def __init__(self, major, minor, ce_pin, irq_pin, temperature, db, sock):
        self.temperature = temperature
        self.db = db
        self.sock = sock
        self.radio = nrf24.NRF24()
        self.radio.begin(major, minor, ce_pin, irq_pin)
        self.radio.setDataRate(self.radio.BR_250KBPS)
        self.radio.setChannel(0)
        self.radio.setAutoAck(1)
        self.radio.enableDynamicPayloads()
        self.radio.printDetails()
        self.radio.openWritingPipe(PIPES[0])
        self.radio.openReadingPipe(1, PIPES[1])

    def run(self):
        try:
            while True:
                self.radio.startListening()
                recv_buffer = self.recv(10)
                self.radio.stopListening()
                if recv_buffer and len(recv_buffer) == 2:
                    self.db.write(1, self.temperature.calc_temp(recv_buffer))
                self.db.write(0, self.temperature.read())
                try:
                    conn, addr = sock.accept()
                except socket.error as e:
                    if e.errno != errno.EAGAIN:
                        raise
                else:
                    try:
                        conn.settimeout(10)
                        recv_line = conn.recv(1024)
                        state, pin = recv_line[:-1].split()
                        cmd = int(pin) << 1 | (state.lower() == 'on')
                        result = self.radio.write(chr(cmd))
                        conn.sendall('%s\n' % ('OK' if result else 'timed out'))
                        print
                        print 'OK' if result else 'timed out'
                    except Exception as e:
                        print
                        print "got invalid line:", repr(recv_line), e
                        try:
                            conn.sendall('invalid request\n')
                        except socket.error:
                            pass
                    finally:
                        conn.close()
        except KeyboardInterrupt:
            print

    def recv(self, timeout=None):
        end = time.time() + timeout
        pipe = [0]
        while not self.radio.available(pipe) and (timeout is None or time.time() < end):
            time.sleep(10000 / 1e6)
        if self.radio.available(pipe):
            recv_buffer = []
            self.radio.read(recv_buffer)
            return recv_buffer
        return []

    def cleanup(self):
        self.radio.end()
        self.db.close()
        self.temperature.cleanup()
        self.sock.close()

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.cleanup()

class DBWriter:
    def __init__(self):
        self.conn = sqlite3.connect('/var/lib/autoboiler/autoboiler.sqlite3')
        self.conn.isolation_level = None
        self.c = self.conn.cursor()
        self.c.execute('''CREATE TABLE IF NOT EXISTS temperature
                          (date datetime, sensor integer, temperature real)''')
        self.c.execute('''CREATE INDEX IF NOT EXISTS temperature_date
                          ON temperature(date)''')

    def write(self, idx, value):
        line = "%s %d %f" % (datetime.datetime.now(), idx, value)
        if idx > 0:
            print '\033[%dC' % len(line * idx),
        print line, '\r',
        sys.stdout.flush()
        try:
            self.c.execute('''insert into temperature values (?, ?, ?)''',
                           (datetime.datetime.now(), idx, value))
        except sqlite3.OperationalError:
            pass

    def close(self):
        self.conn.commit()
        self.c.close()
        self.conn.close()

if __name__ == '__main__':
    GPIO.setmode(GPIO.BCM)
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', required=True, choices=['boiler', 'controller'])
    parser.add_argument('--pidfile',  '-p')
    args = parser.parse_args()
    with open('/etc/default/templog') as f:
        i = int(f.read())
    if args.pidfile:
        with open(args.pidfile, 'w') as f:
            print >>f, os.getpid()
    try:
        if args.mode == 'boiler':
            with Boiler(0, 0, 25, 24, Temperature(0, 1), Relay({0: 17, 1: 18})) as radio:
                radio.run()
        elif args.mode == 'controller':
            sockname = '/var/lib/autoboiler/autoboiler.socket'
            try:
                os.unlink(sockname)
            except OSError as e:
                if e.errno != errno.ENOENT and os.path.exists(sockname):
                    raise
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            sock.bind(sockname)
            os.chmod(sockname, 0777)
            sock.setblocking(0)
            sock.listen(1)
            with Controller(0, 1, 25, 24, Temperature(0, 0), DBWriter(), sock) as radio:
                radio.run()
    finally:
        GPIO.cleanup()
        if args.pidfile:
            os.unlink(args.pidfile)
    sys.exit(0)

# vim: set et sw=4 ts=4 sts=4 ai:
<commit_msg>Add query command to protocol.<commit_after>#!/usr/bin/python

import RPi.GPIO as GPIO
import nrf24
import sys
import time
import spidev
import argparse
import os
import sqlite3
import datetime
import errno
import socket

PIPES = ([0xe7, 0xe7, 0xe7, 0xe7, 0xe7], [0xc2, 0xc2, 0xc2, 0xc2, 0xc2])

class Relay:
    def __init__(self, pins):
        self.pins = pins
        self.states = []
        for pin in self.pins.values():
            GPIO.setup(pin, GPIO.OUT, initial=GPIO.HIGH)
            self.states[pin] = 0

    def output(self, pin, state):
        print "setting pin", pin, state and "on" or "off"
        self.states[pin] = state
        GPIO.output(self.pins[pin], not state)  # These devices are active-low.

    def state(self, pin):
        return self.states[pin]

    def cleanup(self):
        pass  # this will be done later: GPIO.cleanup()

class Temperature:
    def __init__(self, major=0, minor=0):
        self.spi = spidev.SpiDev()
        self.spi.open(major, minor)

    def rawread(self):
        return self.spi.xfer2([0, 0])

    def read(self):
        return self.calc_temp(self.rawread())

    @staticmethod
    def calc_temp(buf):
        return (((buf[0] << 8) | buf[1]) >> 3) * 0.0625

    def cleanup(self):
        self.spi.close()

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.cleanup()

class Boiler:
    def __init__(self, major, minor, ce_pin, irq_pin, temperature, relay):
        self.relay = relay
        self.temperature = temperature
        self.radio = nrf24.NRF24()
        self.radio.begin(major, minor, ce_pin, irq_pin)
        self.radio.setDataRate(self.radio.BR_250KBPS)
        self.radio.setChannel(0)
        self.radio.setAutoAck(1)
        self.radio.enableDynamicPayloads()
        self.radio.printDetails()
        self.radio.openWritingPipe(PIPES[0])
        self.radio.openReadingPipe(1, PIPES[1])

    def run(self):
        while True:
            try:
                self.radio.startListening()
                recv_buffer = self.recv(10)
                self.radio.stopListening()
                for byte in recv_buffer:
                    pin = byte >> 2
                    query = byte >> 1 & 1
                    state = byte & 1
                    if query:
                        self.radio.write(self.relay.state(pin))
                    else:
                        self.relay.output(pin, state)
                result = self.radio.write(self.temperature.rawread())
                if not result:
                    print datetime.datetime.now(), "Did not receive ACK from controller."
            except Exception as e:
                print e

    def recv(self, timeout=None):
        end = time.time() + timeout
        pipe = [0]
        while not self.radio.available(pipe) and (timeout is None or time.time() < end):
            time.sleep(10000 / 1e6)
        if self.radio.available(pipe):
            recv_buffer = []
            self.radio.read(recv_buffer)
            return recv_buffer
        return []

    def cleanup(self):
        self.radio.end()

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.cleanup()

class Controller:
    def __init__(self, major, minor, ce_pin, irq_pin, temperature, db, sock):
        self.temperature = temperature
        self.db = db
        self.sock = sock
        self.radio = nrf24.NRF24()
        self.radio.begin(major, minor, ce_pin, irq_pin)
        self.radio.setDataRate(self.radio.BR_250KBPS)
        self.radio.setChannel(0)
        self.radio.setAutoAck(1)
        self.radio.enableDynamicPayloads()
        self.radio.printDetails()
        self.radio.openWritingPipe(PIPES[0])
        self.radio.openReadingPipe(1, PIPES[1])

    def run(self):
        try:
            while True:
                self.radio.startListening()
                recv_buffer = self.recv(10)
                self.radio.stopListening()
                if recv_buffer and len(recv_buffer) == 2:
                    self.db.write(1, self.temperature.calc_temp(recv_buffer))
                self.db.write(0, self.temperature.read())
                try:
                    conn, addr = sock.accept()
                except socket.error as e:
                    if e.errno != errno.EAGAIN:
                        raise
                else:
                    try:
                        conn.settimeout(10)
                        recv_line = conn.recv(1024)
                        state, pin = recv_line[:-1].split()
                        cmd = int(pin) << 2 | (state.lower() == 'query') << 1 | (state.lower() == 'on')
                        result = self.radio.write(chr(cmd))
                        recv_buffer = []
                        if state.lower() == 'query':
                            self.radio.startListening()
                            recv_buffer = self.recv(1)
                            self.radio.stopListening()
                        if not recv_buffer:
                            recv_buffer = '?'
                        elif len(recv_buffer) == 1:
                            recv_buffer = recv_buffer[0]
                        conn.sendall('%s %s\n' % ('OK' if result else 'timed out', recv_buffer))
                        print
                        print 'OK' if result else 'timed out', recv_line, recv_buffer
                    except Exception as e:
                        print
                        print "got invalid line:", repr(recv_line), e
                        try:
                            conn.sendall('invalid request\n')
                        except socket.error:
                            pass
                    finally:
                        conn.close()
        except KeyboardInterrupt:
            print

    def recv(self, timeout=None):
        end = time.time() + (timeout or 0.0)
        pipe = [0]
        while not self.radio.available(pipe) and (timeout is None or time.time() < end):
            time.sleep(10000 / 1e6)
        if self.radio.available(pipe):
            recv_buffer = []
            self.radio.read(recv_buffer)
            return recv_buffer
        return []

    def cleanup(self):
        self.radio.end()
        self.db.close()
        self.temperature.cleanup()
        self.sock.close()

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.cleanup()

class DBWriter:
    def __init__(self):
        self.conn = sqlite3.connect('/var/lib/autoboiler/autoboiler.sqlite3')
        self.conn.isolation_level = None
        self.c = self.conn.cursor()
        self.c.execute('''CREATE TABLE IF NOT EXISTS temperature
                          (date datetime, sensor integer, temperature real)''')
        self.c.execute('''CREATE INDEX IF NOT EXISTS temperature_date
                          ON temperature(date)''')

    def write(self, idx, value):
        line = "%s %d %f" % (datetime.datetime.now(), idx, value)
        if idx > 0:
            print '\033[%dC' % len(line * idx),
        print line, '\r',
        sys.stdout.flush()
        try:
            self.c.execute('''insert into temperature values (?, ?, ?)''',
                           (datetime.datetime.now(), idx, value))
        except sqlite3.OperationalError:
            pass

    def close(self):
        self.conn.commit()
        self.c.close()
        self.conn.close()

if __name__ == '__main__':
    GPIO.setmode(GPIO.BCM)
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', required=True, choices=['boiler', 'controller'])
    parser.add_argument('--pidfile',  '-p')
    args = parser.parse_args()
    with open('/etc/default/templog') as f:
        i = int(f.read())
    if args.pidfile:
        with open(args.pidfile, 'w') as f:
            print >>f, os.getpid()
    try:
        if args.mode == 'boiler':
            with Boiler(0, 0, 25, 24, Temperature(0, 1), Relay({0: 17, 1: 18})) as radio:
                radio.run()
        elif args.mode == 'controller':
            sockname = '/var/lib/autoboiler/autoboiler.socket'
            try:
                os.unlink(sockname)
            except OSError as e:
                if e.errno != errno.ENOENT and os.path.exists(sockname):
                    raise
            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            sock.bind(sockname)
            os.chmod(sockname, 0777)
            sock.setblocking(0)
            sock.listen(1)
            with Controller(0, 1, 25, 24, Temperature(0, 0), DBWriter(), sock) as radio:
                radio.run()
    finally:
        GPIO.cleanup()
        if args.pidfile:
            os.unlink(args.pidfile)
    sys.exit(0)

# vim: set et sw=4 ts=4 sts=4 ai:
<|endoftext|>