<commit_before>import discord
import os
import json
import subprocess
from datetime import datetime

greetings = "Arcoex stands for **ar**bitrary **co**de **ex**ecution, and that's what this bot does! There are two ways to use me (you can use both at the same time if you want to):\n\n**1.** Mention me among with the language of the code you want to run, and then add the code warped in between \`\`\`, making a code block. (This method overrides the next one if both are specified at the same time) Here's an example:\n\n<@350327901788569612> bash ```\n echo(\"I'm using the first method!\")```\n**2.** Mention me among your code warped in between \`\`\`, and specify the language of your code right after the first set of \`\`\` (This method will add syntax markdown to your code). Here's an example:\n\n<@350327901788569612> ```bash\n echo(\"I'm using the second method!\")```\nSo far, the supported languages are bash, rust, cpp, javascript, guile and python. Message Kurolox if you want to add more!"

def look_for_code(msg):
    """Checks the discord message for a code block. If there's no code, it returns an empty string."""

    try:
        code = msg.split("```")[1] # Looks for anything wrapped in a ``` block
        return code.split("\n")
    except IndexError: # There's no ``` in the message content
        return ""

def detect_language(msg, code):
    """Checks code language and gives the following functions information regarding how to run the code."""

    lang_dict = {}
    for language in os.listdir(os.path.dirname(os.path.realpath(__file__)) + "/languages"):
        with open(os.path.dirname(os.path.realpath(__file__)) + "/languages/" + language) as langjson:
            lang_dict[language[:-5]] = json.load(langjson)

    # Check markdown
    for language in lang_dict.keys():
        if code[0].lower() in lang_dict[language]["known_naming"]:
            code[0] = "" # Erase markdown language
            markdown_language = language

    # Check argument
    try:
        argument_provided = msg.split("```")[0].split(" ")[1]
        for language in lang_dict.keys():
            if argument_provided.lower() in lang_dict[language]["known_naming"]:
                return lang_dict[language], code
    except IndexError:
        pass

    try:
        return lang_dict[markdown_language], code
    except UnboundLocalError: # There's no valid markdown
        return "", code

def create_file(language, extension, code, compiled):
    """create a file with the code provided, so it can be run or compiled later."""

    if not os.path.isdir(os.path.dirname(os.path.realpath(__file__)) + "/code"): # Generates /code folder if it doesn't exist
        os.mkdir(os.path.dirname(os.path.realpath(__file__)) + "/code")

    if not os.path.isdir(os.path.dirname(os.path.realpath(__file__)) + "/code/" + language): # Generates /code/language folder if it doesn't exist
        os.mkdir(os.path.dirname(os.path.realpath(__file__)) + "/code/" + language)

    file_path = "%s/code/%s/%s%s" % (os.path.dirname(os.path.realpath(__file__)), language, datetime.now().strftime("%d-%m-%Y_%H:%M:%S.%f"),  extension)
    with open(file_path, "a") as file: # Open the file to write on
        for line in code:
            file.write("%s\n" % line)

    if compiled: # Creates a special character free copy for compiling if it's needed
        file_compiled = "%s/code/%s/exec" % (os.path.dirname(os.path.realpath(__file__)), language)
        if not os.path.exists(file_compiled):
            os.mkdir(file_compiled)

        try:
            os.remove(file_compiled + "/plaincode" + extension)
        except FileNotFoundError:
            pass
        with open(file_compiled + "/plaincode" + extension, "a") as file: # Open the file to write on
            for line in code:
                file.write("%s\n" % line)
    else:
        file_compiled = ""

    return file_path, file_compiled

def run_compiler(file_path, language, compiler_exec, extension, flags = ""):
    """runs the compiler and generates an executable."""

    run_process = subprocess.Popen(["umlbox", "--cwd", "/arcoex", "-B", "-fw", "/arcoex", "-fw", "/bot/code", "-m", "512M", "-T", "65", compiler_exec] + flags + [file_path + "/executable", file_path + "/plaincode" + extension], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    try:
        stdout, stderr = run_process.communicate(timeout=60)
        if stdout:
            return "", stdout.decode("utf-8")
        elif stderr:
            return "", stderr.decode("utf-8")
        else:
            return file_path + "/executable", ""
    except subprocess.TimeoutExpired:
        run_process.kill()
        return 0

def execute_code(file_path, run_command, flags):
    """Creates a subprocess that runs the file, then grabs the stdout and stderr of said process."""

    timeout_flag = False
    if run_command: # For non-compiled languages
        run_process = subprocess.Popen(["umlbox", "--cwd", "/arcoex", "-B", "-fw", "/arcoex", "-f", "/bot/code", "-m", "512M", "-T", "10", run_command] + flags + [file_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    else:
        run_process = subprocess.Popen(["umlbox", "--cwd", "/arcoex", "-B", "-fw", "/arcoex", "-f", "/bot/code", "-m", "512M", "-T", "10", file_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        stdout, stderr = run_process.communicate(timeout=8)
    except subprocess.TimeoutExpired:
        stdout, stderr, timeout_flag = timeout(run_process)

    return stdout.decode("utf-8"), stderr.decode("utf-8"), timeout_flag

def timeout(process):
    """If the process timed out, returns the last 5 lines of output and kills the process"""
    line_number = 0
    stderr = b""
    stdout = b""
    while line_number < 5:
        stdout += process.stdout.readline()
        line_number += 1
    process.kill()
    timeout_flag = True
    return stdout, stderr, timeout_flag

async def bot_reply(msg, client, signal, code_output="", compiler_output="", language=""):
    """Makes the bot reply depending on the signal."""

    if signal == 0: # There was no code detected on the message
        await client.send_message(msg.channel, "I'm sorry, but I'm not seeing any code to run.")
    elif signal == 1: # There was only a mention of the bot
        await client.send_message(msg.channel, greetings)
    elif signal == 2: # detect_language() couldn't detect the language
        await client.send_message(msg.channel, "I'm sorry, but you didn't specify the language of the code or it isn't supported yet.")
    elif signal == 3: # The compiler ran into issues
        await client.send_message(msg.channel, "I've encountered the following error when compiling your %s code. \n```\n%s```" % (language, compiler_output))
    elif signal == 4: # The process timeouted
        await client.send_message(msg.channel, "I'm sorry, but your request took too long. Here are the last 5 lines from the output.\n```\n%s```" % code_output)
    elif signal == 5: # There were errors when running the code
        await client.send_message(msg.channel, "I've encountered the following error when running your %s code. \n```\n%s```" % (language, code_output))
    elif signal == 6: # Everything went well
        await client.send_message(msg.channel, "Here's the output of your %s code. \n```\n%s```" % (language, code_output))

async def code(msg, client):
    """Grab the code block from a message, run it and return the output."""

    code = look_for_code(msg.content)
    if not code:
        if msg.content != ("<@%s>" % client.user.id): # if the message content is not just a mention
            await bot_reply(msg, client, 0)
        else: # if the message content is only a mention to the bot
            await bot_reply(msg, client, 1)
    else:
        lang_json, code = detect_language(msg.content, code)
        if not lang_json: # The language wasn't detected
            await bot_reply(msg, client, 2)
        else:
            file_path, compiled_copy = create_file(lang_json["language"], lang_json["file_extension"], code, lang_json["compiled"])
            if lang_json["compiled"]: # Code needs to be compiled
                exec_path, compiler_output = run_compiler(compiled_copy, lang_json["language"], lang_json["compiler_exec"],lang_json["file_extension"], flags= lang_json["compiler_flags"])
                if compiler_output: # Errors when compiling
                    await bot_reply(msg, client, 3, compiler_output=compiler_output, language=lang_json["language"])
                else:
                    output, error, timeout_flag = execute_code(exec_path, lang_json["exec"], lang_json["exec_flags"])
                    if timeout_flag: # The process did timeout
                        await bot_reply(msg, client, 4, code_output=output)
                    elif error: # There's something on stderr
                        await bot_reply(msg, client, 5, code_output=error, language=lang_json["language"])
                    else: # Everything went fine
                        await bot_reply(msg, client, 6, code_output=output, language=lang_json["language"])

            else: # Code can be run directly
                output, error, timeout_flag = execute_code(file_path, lang_json["exec"], lang_json["exec_flags"])
                if timeout_flag: # The process did timeout
                    await bot_reply(msg, client, 4, code_output=output)
                elif error: # There's something on stderr
                    await bot_reply(msg, client, 5, code_output=error, language=lang_json["language"])
                else: # Everything went fine
                    await bot_reply(msg, client, 6, code_output=output, language=lang_json["language"])
<commit_msg>Added support for hastebin<commit_after>import discord
import os
import json
import subprocess
import hastebin
from datetime import datetime

greetings = "Arcoex stands for **ar**bitrary **co**de **ex**ecution, and that's what this bot does! There are two ways to use me (you can use both at the same time if you want to):\n\n**1.** Mention me among with the language of the code you want to run, and then add the code warped in between \`\`\`, making a code block. (This method overrides the next one if both are specified at the same time) Here's an example:\n\n<@350327901788569612> bash ```\n echo(\"I'm using the first method!\")```\n**2.** Mention me among your code warped in between \`\`\`, and specify the language of your code right after the first set of \`\`\` (This method will add syntax markdown to your code). Here's an example:\n\n<@350327901788569612> ```bash\n echo(\"I'm using the second method!\")```\nSo far, the supported languages are bash, rust, cpp, javascript, guile and python. Message Kurolox if you want to add more!"

def look_for_code(msg):
    """Checks the discord message for a code block. If there's no code, it returns an empty string."""

    try:
        code = msg.split("```")[1] # Looks for anything wrapped in a ``` block
        return code.split("\n")
    except IndexError: # There's no ``` in the message content
        return ""

def detect_language(msg, code):
    """Checks code language and gives the following functions information regarding how to run the code."""

    lang_dict = {}
    for language in os.listdir(os.path.dirname(os.path.realpath(__file__)) + "/languages"):
        with open(os.path.dirname(os.path.realpath(__file__)) + "/languages/" + language) as langjson:
            lang_dict[language[:-5]] = json.load(langjson)

    # Check markdown
    for language in lang_dict.keys():
        if code[0].lower() in lang_dict[language]["known_naming"]:
            code[0] = "" # Erase markdown language
            markdown_language = language

    # Check argument
    try:
        argument_provided = msg.split("```")[0].split(" ")[1]
        for language in lang_dict.keys():
            if argument_provided.lower() in lang_dict[language]["known_naming"]:
                return lang_dict[language], code
    except IndexError:
        pass

    try:
        return lang_dict[markdown_language], code
    except UnboundLocalError: # There's no valid markdown
        return "", code

def create_file(language, extension, code, compiled):
    """create a file with the code provided, so it can be run or compiled later."""

    if not os.path.isdir(os.path.dirname(os.path.realpath(__file__)) + "/code"): # Generates /code folder if it doesn't exist
        os.mkdir(os.path.dirname(os.path.realpath(__file__)) + "/code")

    if not os.path.isdir(os.path.dirname(os.path.realpath(__file__)) + "/code/" + language): # Generates /code/language folder if it doesn't exist
        os.mkdir(os.path.dirname(os.path.realpath(__file__)) + "/code/" + language)

    file_path = "%s/code/%s/%s%s" % (os.path.dirname(os.path.realpath(__file__)), language, datetime.now().strftime("%d-%m-%Y_%H:%M:%S.%f"),  extension)
    with open(file_path, "a") as file: # Open the file to write on
        for line in code:
            file.write("%s\n" % line)

    if compiled: # Creates a special character free copy for compiling if it's needed
        file_compiled = "%s/code/%s/exec" % (os.path.dirname(os.path.realpath(__file__)), language)
        if not os.path.exists(file_compiled):
            os.mkdir(file_compiled)

        try:
            os.remove(file_compiled + "/plaincode" + extension)
        except FileNotFoundError:
            pass
        with open(file_compiled + "/plaincode" + extension, "a") as file: # Open the file to write on
            for line in code:
                file.write("%s\n" % line)
    else:
        file_compiled = ""

    return file_path, file_compiled

def run_compiler(file_path, language, compiler_exec, extension, flags = ""):
    """runs the compiler and generates an executable."""

    run_process = subprocess.Popen(["umlbox", "--cwd", "/arcoex", "-B", "-fw", "/arcoex", "-fw", "/bot/code", "-m", "512M", "-T", "65", compiler_exec] + flags + [file_path + "/executable", file_path + "/plaincode" + extension], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    try:
        stdout, stderr = run_process.communicate(timeout=60)
        if stdout:
            return "", stdout.decode("utf-8")
        elif stderr:
            return "", stderr.decode("utf-8")
        else:
            return file_path + "/executable", ""
    except subprocess.TimeoutExpired:
        run_process.kill()
        return 0

def execute_code(file_path, run_command, flags):
    """Creates a subprocess that runs the file, then grabs the stdout and stderr of said process."""

    timeout_flag = False
    if run_command: # For non-compiled languages
        run_process = subprocess.Popen(["umlbox", "--cwd", "/arcoex", "-B", "-fw", "/arcoex", "-f", "/bot/code", "-m", "512M", "-T", "10", run_command] + flags + [file_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    else:
        run_process = subprocess.Popen(["umlbox", "--cwd", "/arcoex", "-B", "-fw", "/arcoex", "-f", "/bot/code", "-m", "512M", "-T", "10", file_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        stdout, stderr = run_process.communicate(timeout=8)
    except subprocess.TimeoutExpired:
        stdout, stderr, timeout_flag = timeout(run_process)

    return stdout.decode("utf-8"), stderr.decode("utf-8"), timeout_flag

def timeout(process):
    """If the process timed out, returns the last 5 lines of output and kills the process"""
    line_number = 0
    stderr = b""
    stdout = b""
    while line_number < 5:
        stdout += process.stdout.readline()
        line_number += 1
    process.kill()
    timeout_flag = True
    return stdout, stderr, timeout_flag

async def bot_reply(msg, client, signal, code_output="", compiler_output="", language=""):
    """Makes the bot reply depending on the signal."""

    if signal == 0: # There was no code detected on the message
        await client.send_message(msg.channel, "I'm sorry, but I'm not seeing any code to run.")
    elif signal == 1: # There was only a mention of the bot
        await client.send_message(msg.channel, greetings)
    elif signal == 2: # detect_language() couldn't detect the language
        await client.send_message(msg.channel, "I'm sorry, but you didn't specify the language of the code or it isn't supported yet.")
    elif signal == 3: # The compiler ran into issues
        await client.send_message(msg.channel, "I've encountered the following error when compiling your %s code. \n```\n%s```" % (language, compiler_output))
    elif signal == 4: # The process timeouted
        await client.send_message(msg.channel, "I'm sorry, but your request took too long. Here are the last 5 lines from the output.\n```\n%s```" % code_output)
    elif signal == 5: # There were errors when running the code
        await client.send_message(msg.channel, "I've encountered the following error when running your %s code. \n```\n%s```" % (language, code_output))
    elif signal == 6: # Everything went well
        try:
            await client.send_message(msg.channel, "Here's the output of your %s code. \n```\n%s```" % (language, code_output))
        except discord.errors.HTTPException: # Unless the message was +2000 characters
            await client.send_message(msg.channel, "I'm sorry, but the output is too long for Discord. You can check the output of your %s code here. %s" % (language, hastebin.post(code_output)))

async def code(msg, client):
    """Grab the code block from a message, run it and return the output."""

    code = look_for_code(msg.content)
    if not code:
        if msg.content != ("<@%s>" % client.user.id): # if the message content is not just a mention
            await bot_reply(msg, client, 0)
        else: # if the message content is only a mention to the bot
            await bot_reply(msg, client, 1)
    else:
        lang_json, code = detect_language(msg.content, code)
        if not lang_json: # The language wasn't detected
            await bot_reply(msg, client, 2)
        else:
            file_path, compiled_copy = create_file(lang_json["language"], lang_json["file_extension"], code, lang_json["compiled"])
            if lang_json["compiled"]: # Code needs to be compiled
                exec_path, compiler_output = run_compiler(compiled_copy, lang_json["language"], lang_json["compiler_exec"],lang_json["file_extension"], flags= lang_json["compiler_flags"])
                if compiler_output: # Errors when compiling
                    await bot_reply(msg, client, 3, compiler_output=compiler_output, language=lang_json["language"])
                else:
                    output, error, timeout_flag = execute_code(exec_path, lang_json["exec"], lang_json["exec_flags"])
                    if timeout_flag: # The process did timeout
                        await bot_reply(msg, client, 4, code_output=output)
                    elif error: # There's something on stderr
                        await bot_reply(msg, client, 5, code_output=error, language=lang_json["language"])
                    else: # Everything went fine
                        await bot_reply(msg, client, 6, code_output=output, language=lang_json["language"])

            else: # Code can be run directly
                output, error, timeout_flag = execute_code(file_path, lang_json["exec"], lang_json["exec_flags"])
                if timeout_flag: # The process did timeout
                    await bot_reply(msg, client, 4, code_output=output)
                elif error: # There's something on stderr
                    await bot_reply(msg, client, 5, code_output=error, language=lang_json["language"])
                else: # Everything went fine
                    await bot_reply(msg, client, 6, code_output=output, language=lang_json["language"])
<|endoftext|>


<commit_before>import logging
import socket
import threading

from channel import Channel, ChannelList

class Client:
	def __init__(self, server, nick, username, realname, **conf):
		self.sock = None
		self.nick = None
		self.server = None

		self.conf = {
			"port": 6667,
			"autorejoin": True,
			"autoconn": True,
			"channels": [],
			"prefixes": "&#"
		}

		self.conf.update(conf)
		self.conf.update({
			"server": server,
			"nick": nick,
			"username": username,
			"realname": realname,
		})

		self.channels = ChannelList()

		if self.conf["autoconn"]:
			thread = threading.Thread(name=self.conf["server"], target=self.connect)
			logging.debug("Dispatching server thread %s", thread.name)
			thread.start()

	def connect(self):
		"""Attempt to connect to the server. Log in if successful."""

		self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

		try:
			self.sock.connect((self.conf["server"], self.conf["port"]))
			self.server = self.conf["server"]
			logging.info("%s connected", self.server)

			self.send("NICK", self.conf["nick"])
			self.send("USER", self.conf["username"], 8, "*", self.conf["realname"])

			thread = threading.Thread(name=self.conf["server"] + "-listen", target=self.listen)
			logging.debug("Dispatching listener thread %s", thread.name)
			thread.start()
		except socket.error as e:
			logging.error("%s cannot connect: %s", self.conf["server"], e)
			self.sock.close()
			self.sock = None

	def send(self, *msg):
		"""Send a raw message to the server. Prepend a colon to the last parameter."""

		if not self.sock:
			logging.error("%s not connected to server", self.conf["server"])
			return

		message = " ".join(map(str, msg[:-1])) + " :" + str(msg[-1])
		self.sock.send(bytes(message + "\r\n", "utf-8"))
		logging.debug("Send: %s", message.rstrip())

	def listen(self):
		"""Listen for incoming messages from the IRC server."""
		pass
<commit_msg>use sendall instead of send<commit_after>import logging
import socket
import threading

from channel import Channel, ChannelList

class Client:
	def __init__(self, server, nick, username, realname, **conf):
		self.sock = None
		self.nick = None
		self.server = None

		self.conf = {
			"port": 6667,
			"autorejoin": True,
			"autoconn": True,
			"channels": [],
			"prefixes": "&#"
		}

		self.conf.update(conf)
		self.conf.update({
			"server": server,
			"nick": nick,
			"username": username,
			"realname": realname,
		})

		self.channels = ChannelList()

		if self.conf["autoconn"]:
			thread = threading.Thread(name=self.conf["server"], target=self.connect)
			logging.debug("Dispatching server thread %s", thread.name)
			thread.start()

	def connect(self):
		"""Attempt to connect to the server. Log in if successful."""

		self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

		try:
			self.sock.connect((self.conf["server"], self.conf["port"]))
			self.server = self.conf["server"]
			logging.info("%s connected", self.server)

			self.send("NICK", self.conf["nick"])
			self.send("USER", self.conf["username"], 8, "*", self.conf["realname"])

			thread = threading.Thread(name=self.conf["server"] + "-listen", target=self.listen)
			logging.debug("Dispatching listener thread %s", thread.name)
			thread.start()
		except socket.error as e:
			logging.error("%s cannot connect: %s", self.conf["server"], e)
			self.sock.close()
			self.sock = None

	def send(self, *msg):
		"""Send a raw message to the server. Prepend a colon to the last parameter."""

		if not self.sock:
			logging.error("%s not connected to server", self.conf["server"])
			return

		message = " ".join(map(str, msg[:-1])) + " :" + str(msg[-1])
		self.sock.sendall(bytes(message + "\r\n", "utf-8"))
		logging.debug("Send: %s", message.rstrip())

	def listen(self):
		"""Listen for incoming messages from the IRC server."""
		pass
<|endoftext|>


<commit_before>#
# Copyright (C) 2007-2013 by Johan De Taeye, frePPLe bvba
#
# This library is free software; you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero
# General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
from __future__ import print_function

from django.db import models, transaction, DEFAULT_DB_ALIAS
from django.db.models.signals import post_save
from django.dispatch import receiver
from django.utils.translation import ugettext_lazy as _
from django.conf import settings

from freppledb.common.models import User

import logging
logger = logging.getLogger(__name__)

class Task(models.Model):
  '''
  Expected status values are:
    - 'Waiting'
    - 'Done'
    - 'Failed'
    - 'Canceled'
    - 'DD%', where DD represents the percentage completed
  Other values are okay, but the above have translations.
  '''
  # Database fields
  id = models.AutoField(_('identifier'), primary_key=True, editable=False)
  name = models.CharField(_('name'), max_length=20, db_index=True, editable=False)
  submitted = models.DateTimeField(_('submitted'), editable=False)
  started = models.DateTimeField(_('started'), blank=True, null=True, editable=False)
  finished = models.DateTimeField(_('submitted'), blank=True, null=True, editable=False)
  arguments = models.TextField(_('arguments'), max_length=200, null=True, editable=False)
  status = models.CharField(_('status'), max_length=20, editable=False)
  message = models.TextField(_('message'), max_length=200, null=True, editable=False)
  user = models.ForeignKey(User, verbose_name=_('user'), blank=True, null=True, editable=False)

  def __unicode__(self):
    return self.id + ' - ' + self.name + ' - ' + self.status

  class Meta:
    db_table = "execute_log"
    verbose_name_plural = _('tasks')
    verbose_name = _('task')

  @staticmethod
  def submitTask():
    # Add record to the database
    # Check if a worker is present. If not launch one.
    return 1

class Scenario(models.Model):
  scenarioStatus = (
    ('free',_('Free')),
    ('in use',_('In use')),
    ('busy',_('Busy')),
  )

  # Database fields
  name = models.CharField(_('name'), max_length=settings.NAMESIZE, primary_key=True)
  description = models.CharField(_('description'), max_length=settings.DESCRIPTIONSIZE, null=True, blank=True)
  status = models.CharField(_('status'), max_length=10,
    null=False, blank=False, choices=scenarioStatus
    )
  lastrefresh = models.DateTimeField(_('last refreshed'), null=True, editable=False)

  def __unicode__(self):
    return self.name

  @staticmethod
  def syncWithSettings():
    ac = transaction.get_autocommit()
    transaction.set_autocommit(False)
    try:
      # Bring the scenario table in sync with settings.databases
      dbs = [ i for i,j in settings.DATABASES.items() if j['NAME'] ]
      for sc in Scenario.objects.all():
        if sc.name not in dbs:
          sc.delete()
      scs = [sc.name for sc in Scenario.objects.all()]
      for db in dbs:
        if db not in scs:
          if db == DEFAULT_DB_ALIAS:
            Scenario(name=db, status=u"In use", description='Production database').save()
          else:
            Scenario(name=db, status=u"Free").save()
      transaction.commit()
    except Exception as e:
      logger.error("Error synchronizing the scenario table with the settings: %s" % e)
      transaction.rollback()
    finally:
      transaction.set_autocommit(ac)

  class Meta:
    db_table = "execute_scenario"
    permissions = (
        ("copy_scenario", "Can copy a scenario"),
        ("release_scenario", "Can release a scenario"),
       )
    verbose_name_plural = _('scenarios')
    verbose_name = _('scenario')
    ordering = ['name']

@receiver(post_save, sender=User)
def sync_handler(sender, **kwargs):
  if not kwargs.get('created',False) or kwargs.get('using',DEFAULT_DB_ALIAS) != DEFAULT_DB_ALIAS:
    return
  # A new user is created in the default database.
  # We create the same user in all scenarios that are in use. Otherwise the user can't create
  # comments or edit objects in these what-if scenarios.
  for sc in Scenario.objects.all().filter(status=u'In use'):
    kwargs['instance'].save(using=sc.name)
<commit_msg>updated printing of Task model<commit_after>#
# Copyright (C) 2007-2013 by Johan De Taeye, frePPLe bvba
#
# This library is free software; you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero
# General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
from __future__ import print_function

from django.db import models, transaction, DEFAULT_DB_ALIAS
from django.db.models.signals import post_save
from django.dispatch import receiver
from django.utils.translation import ugettext_lazy as _
from django.conf import settings

from freppledb.common.models import User

import logging
logger = logging.getLogger(__name__)

class Task(models.Model):
  '''
  Expected status values are:
    - 'Waiting'
    - 'Done'
    - 'Failed'
    - 'Canceled'
    - 'DD%', where DD represents the percentage completed
  Other values are okay, but the above have translations.
  '''
  # Database fields
  id = models.AutoField(_('identifier'), primary_key=True, editable=False)
  name = models.CharField(_('name'), max_length=20, db_index=True, editable=False)
  submitted = models.DateTimeField(_('submitted'), editable=False)
  started = models.DateTimeField(_('started'), blank=True, null=True, editable=False)
  finished = models.DateTimeField(_('submitted'), blank=True, null=True, editable=False)
  arguments = models.TextField(_('arguments'), max_length=200, null=True, editable=False)
  status = models.CharField(_('status'), max_length=20, editable=False)
  message = models.TextField(_('message'), max_length=200, null=True, editable=False)
  user = models.ForeignKey(User, verbose_name=_('user'), blank=True, null=True, editable=False)

  def __unicode__(self):
    return "%s - %s - %s" % (self.id, self.name, self.status)

  class Meta:
    db_table = "execute_log"
    verbose_name_plural = _('tasks')
    verbose_name = _('task')

  @staticmethod
  def submitTask():
    # Add record to the database
    # Check if a worker is present. If not launch one.
    return 1

class Scenario(models.Model):
  scenarioStatus = (
    ('free',_('Free')),
    ('in use',_('In use')),
    ('busy',_('Busy')),
  )

  # Database fields
  name = models.CharField(_('name'), max_length=settings.NAMESIZE, primary_key=True)
  description = models.CharField(_('description'), max_length=settings.DESCRIPTIONSIZE, null=True, blank=True)
  status = models.CharField(_('status'), max_length=10,
    null=False, blank=False, choices=scenarioStatus
    )
  lastrefresh = models.DateTimeField(_('last refreshed'), null=True, editable=False)

  def __unicode__(self):
    return self.name

  @staticmethod
  def syncWithSettings():
    ac = transaction.get_autocommit()
    transaction.set_autocommit(False)
    try:
      # Bring the scenario table in sync with settings.databases
      dbs = [ i for i,j in settings.DATABASES.items() if j['NAME'] ]
      for sc in Scenario.objects.all():
        if sc.name not in dbs:
          sc.delete()
      scs = [sc.name for sc in Scenario.objects.all()]
      for db in dbs:
        if db not in scs:
          if db == DEFAULT_DB_ALIAS:
            Scenario(name=db, status=u"In use", description='Production database').save()
          else:
            Scenario(name=db, status=u"Free").save()
      transaction.commit()
    except Exception as e:
      logger.error("Error synchronizing the scenario table with the settings: %s" % e)
      transaction.rollback()
    finally:
      transaction.set_autocommit(ac)

  class Meta:
    db_table = "execute_scenario"
    permissions = (
        ("copy_scenario", "Can copy a scenario"),
        ("release_scenario", "Can release a scenario"),
       )
    verbose_name_plural = _('scenarios')
    verbose_name = _('scenario')
    ordering = ['name']

@receiver(post_save, sender=User)
def sync_handler(sender, **kwargs):
  if not kwargs.get('created',False) or kwargs.get('using',DEFAULT_DB_ALIAS) != DEFAULT_DB_ALIAS:
    return
  # A new user is created in the default database.
  # We create the same user in all scenarios that are in use. Otherwise the user can't create
  # comments or edit objects in these what-if scenarios.
  for sc in Scenario.objects.all().filter(status=u'In use'):
    kwargs['instance'].save(using=sc.name)
<|endoftext|>


<commit_before>"""
My purpose in life is to daily check the VTEC database and see if there are
any IDs that are missing.  daryl then follows up with the weather bureau
reporting anything he finds after an investigation.

$Id: $:
"""
import mx.DateTime
from pyIEM import iemdb
i = iemdb.iemdb()
postgis = i['postgis']

tests = ['HGX.FF.7', 'GUM.MA.1', 'GUM.MA.2','MRX.SV.18', 'SGF.TO.23',
         'LMK.TO.81', 'LMK.TO.82', 'LSX.TO.11', 'IND.TO.7', 'IND.TO.8',
         'PAH.TO.52', 'PAH.TO.53','PAH.TO.54', 'PAH.TO.55', 'MRX.TO.58',
         'ILX.TO.2', 'CHS.TO.6', 'CAE.FF.1', 'CAE.TO.4', 'CAE.TO.5',
         'RNK.TO.6', 'ICT.TO.5', 'TOP.TO.5',
]
#for i in range(27,78):
#    tests.append('OAX.TO.%s' % (i,))
sql = "SELECT wfo, min(eventid), max(eventid), phenomena from warnings_%s \
       WHERE phenomena IN ('MA','FF','SV','TO') and significance = 'W' \
       GROUP by wfo, phenomena" % (mx.DateTime.now().year, )
rs = postgis.query(sql).dictresult()
for i in range(len(rs)):
  sEvent = int(rs[i]['min'])
  eEvent = int(rs[i]['max'])
  phenomena = rs[i]['phenomena']
  wfo = rs[i]['wfo']
  for e in range(1,sEvent):
    lookup = "%s.%s.%s" % (wfo, phenomena, e)
    if lookup not in tests:
      print "Warning Missing WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, e)

  for eventid in range(sEvent, eEvent):
    sql = "SELECT gtype, count(*) as c from warnings_%s WHERE wfo = '%s' \
           and phenomena = '%s' and eventid = '%s' and significance = 'W' \
           GROUP by gtype" % (mx.DateTime.now().year, wfo,\
           phenomena, eventid)
    rs2 = postgis.query(sql).dictresult()
    polyCount = 0
    cntyCount = 0
    for q in range(len(rs2)):
      if (rs2[q]['gtype'] == "P"): polyCount = int(rs2[q]['c'])
      if (rs2[q]['gtype'] == "C"): cntyCount = int(rs2[q]['c'])

    if (cntyCount == 0 and polyCount == 0):
      lookup = "%s.%s.%s" % (wfo, phenomena, eventid)
      if lookup not in tests:
        print "Warning Missing WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
    elif (polyCount == 0):
      print "SBW Missing     WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
    elif (cntyCount == 0):
      print "County Missing  WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
    elif (polyCount > 1):
      print "Duplicate SBW   WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
      sql = "DELETE from warnings_%s WHERE oid IN (\
         SELECT max(oid) as m from warnings_%s WHERE wfo = '%s' \
         and phenomena = '%s' and eventid = '%s' and significance = 'W' \
         and gtype = 'P')" % (mx.DateTime.now().year, mx.DateTime.now().year,\
          wfo, phenomena, eventid)
      postgis.query(sql)
      sql = "DELETE from warnings_%s WHERE oid IN (\
              select m from (\
select ugc,  max(oid) as m, count(oid) as c from warnings_%s WHERE wfo = '%s' \
and phenomena = '%s' and eventid = '%s' and significance = 'W' \
and gtype = 'C' GROUP by ugc) as foo WHERE c > 1)" % (mx.DateTime.now().year, mx.DateTime.now().year,\
          wfo, phenomena, eventid)
      postgis.query(sql)
<commit_msg>more knowns<commit_after>"""
My purpose in life is to daily check the VTEC database and see if there are
any IDs that are missing.  daryl then follows up with the weather bureau
reporting anything he finds after an investigation.

$Id: $:
"""
import mx.DateTime
from pyIEM import iemdb
i = iemdb.iemdb()
postgis = i['postgis']

tests = ['HGX.FF.7', 'GUM.MA.1', 'GUM.MA.2','MRX.SV.18', 'SGF.TO.23',
         'LMK.TO.81', 'LMK.TO.82', 'LSX.TO.11', 'IND.TO.7', 'IND.TO.8',
         'PAH.TO.52', 'PAH.TO.53','PAH.TO.54', 'PAH.TO.55', 'MRX.TO.58',
         'ILX.TO.2', 'CHS.TO.6', 'CAE.FF.1', 'CAE.TO.4', 'CAE.TO.5',
         'RNK.TO.6', 'ICT.TO.5', 'TOP.TO.5','EAX.TO.1', 'OAX.SV.22',
         'OAX.SV.21', 'OAX.SV.11', 'OAX.SV.13', 'OAX.SV.16', 'OAX.SV.17'
]
#for i in range(27,78):
#    tests.append('OAX.TO.%s' % (i,))
sql = "SELECT wfo, min(eventid), max(eventid), phenomena from warnings_%s \
       WHERE phenomena IN ('MA','FF','SV','TO') and significance = 'W' \
       GROUP by wfo, phenomena" % (mx.DateTime.now().year, )
rs = postgis.query(sql).dictresult()
for i in range(len(rs)):
  sEvent = int(rs[i]['min'])
  eEvent = int(rs[i]['max'])
  phenomena = rs[i]['phenomena']
  wfo = rs[i]['wfo']
  for e in range(1,sEvent):
    lookup = "%s.%s.%s" % (wfo, phenomena, e)
    if lookup not in tests:
      print "Warning Missing WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, e)

  for eventid in range(sEvent, eEvent):
    sql = "SELECT gtype, count(*) as c from warnings_%s WHERE wfo = '%s' \
           and phenomena = '%s' and eventid = '%s' and significance = 'W' \
           GROUP by gtype" % (mx.DateTime.now().year, wfo,\
           phenomena, eventid)
    rs2 = postgis.query(sql).dictresult()
    polyCount = 0
    cntyCount = 0
    for q in range(len(rs2)):
      if (rs2[q]['gtype'] == "P"): polyCount = int(rs2[q]['c'])
      if (rs2[q]['gtype'] == "C"): cntyCount = int(rs2[q]['c'])

    if (cntyCount == 0 and polyCount == 0):
      lookup = "%s.%s.%s" % (wfo, phenomena, eventid)
      if lookup not in tests:
        print "Warning Missing WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
    elif (polyCount == 0):
      print "SBW Missing     WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
    elif (cntyCount == 0):
      print "County Missing  WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
    elif (polyCount > 1):
      print "Duplicate SBW   WFO: %s PHENOMENA: %s EVENTID: %s" % (wfo, phenomena, eventid)
      sql = "DELETE from warnings_%s WHERE oid IN (\
         SELECT max(oid) as m from warnings_%s WHERE wfo = '%s' \
         and phenomena = '%s' and eventid = '%s' and significance = 'W' \
         and gtype = 'P')" % (mx.DateTime.now().year, mx.DateTime.now().year,\
          wfo, phenomena, eventid)
      postgis.query(sql)
      sql = "DELETE from warnings_%s WHERE oid IN (\
              select m from (\
select ugc,  max(oid) as m, count(oid) as c from warnings_%s WHERE wfo = '%s' \
and phenomena = '%s' and eventid = '%s' and significance = 'W' \
and gtype = 'C' GROUP by ugc) as foo WHERE c > 1)" % (mx.DateTime.now().year, mx.DateTime.now().year,\
          wfo, phenomena, eventid)
      postgis.query(sql)
<|endoftext|>


<commit_before>#coding:utf-8
from __future__ import absolute_import
from __future__ import print_function
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Convolution2D, MaxPooling2D
from keras.optimizers import SGD
from keras.utils import np_utils, generic_utils
from six.moves import range
from keras.callbacks import EarlyStopping
from keras.callbacks import ModelCheckpoint
from keras.regularizers import l2, activity_l2
from keras.layers.advanced_activations import PReLU
from data_train_valid_01 import load_data
import h5py
import argparse
from keras.models import model_from_json
#import numpy as np
#np.random.seed(1337)  # for reproducibility
#from keras.utils.visualize_util import plot
#import tensorflow as tf
#import theano
#theano.config.device = 'gpu1'
#theano.config.floatX = 'float32'

#Number of samples per gradient update
batch_size = 100
#分类个数
nb_classes = 2
#the number of times to iterate over the training data arrays
nb_epoch = 20
#训练CNN的个数
nb_loop = 50

# input image dimensions
img_rows, img_cols = 64, 32
# the images are RGB
img_channels = 3
'''
ap = argparse.ArgumentParser()
ap.add_argument("--train", required=True, help="Path to the training set")
ap.add_argument("--valid", required=True, help="Path to the validation set")
ap.add_argument("--test", required=True, help="Path to the testing set")
ap.add_argument("--0", required=True, help="Path to the negative set")
args = vars(ap.parse_args())

#加载数据
path_train = args["train"]
path_valid = args["valid"]
path_test = args["test"]
path_0 = args["0"]
'''
#path_train='/home/cad/dataset/middlepink/data/train'
#path_valid='/home/cad/dataset/middlepink/data/valid_test_1'
#path_train_0='/home/cad/dataset/middlepink/train_negative'
#path_valid_0='/home/cad/dataset/middlepink/data/valid_test_0'

path_train='/home/cad/dataset/pinkall/train_1'
path_valid='/home/cad/dataset/pinkall/valid_test_1'
path_train_0='/home/cad/dataset/pinkall/train_0'
path_valid_0='/home/cad/dataset/pinkall/valid_test_0'

X_train,Y_train_label,X_valid,Y_valid_label = load_data(path_train,path_valid,path_train_0,path_valid_0)
print(X_train.shape[0], 'training samples')
print(X_valid.shape[0],'validation samples')
#print(X_test.shape[0],'testing samples')

#label为0,1两个类别，keras要求格式为binary class matrices,转化一下，调用keras提供的函数
Y_train = np_utils.to_categorical(Y_train_label, nb_classes)
Y_valid = np_utils.to_categorical(Y_valid_label, nb_classes)
#Y_test = np_utils.to_categorical(Y_test_label, nb_classes)

dist={}
#prelu=PReLU(init='zero', weights=None)
for i in range(nb_loop):
    ###############
    #开始建立CNN模型
    ###############
    model = Sequential()
    #第一个卷积层，32个卷积核，每个卷积核大小5*5
    #border_mode为valid
    #激活函数用relu
    model.add(Convolution2D(32, 5, 5,
            #        border_mode='valid',
            #        activation='relu',
            #        W_regularizer=l2(0.01),
            #        b_regularizer=l2(0.01),
            #       activity_regularizer=activity_l2(0.001),
                    input_shape=(img_channels, img_rows, img_cols)))
                    #input_shape=(X_train.shape[1:])))
    model.add(Activation('relu'))
#    model.add(prelu)
#    model.add(MaxPooling2D(pool_size=(2, 2)))
#    model.add(Dropout(0.5))

    #第二个卷积层，32个卷积核，每个卷积核大小5*5
    #采用maxpooling,poolsize为（2,2）
    #使用Dropout
    model.add(Convolution2D(32, 5, 5))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.5))

    #第三个卷积层，64个卷积核，每个卷积核大小3*3
    model.add(Convolution2D(64, 3, 3))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.5))

    #全连接层
    model.add(Flatten())
    model.add(Dense(128))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))

    #Softmax回归，输出2个类别
    model.add(Dense(nb_classes))
    model.add(Activation('softmax'))

    ##############
    #开始训练模型
    ##############

    #lr: float >= 0. Learning rate.
    #momentum: float >= 0. Parameter updates momentum.
    #decay: float >= 0. Learning rate decay over each update.
    #nesterov: boolean. Whether to apply Nesterov momentum.
    #model.compile里的参数loss就是损失函数(目标函数),binary_crossentropy即为logistic loss
    sgd = SGD(lr=0.004, decay=1e-6, momentum=0.9, nesterov=True)
    model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])

    #this will do preprocessing and realtime data augmentation
    '''
    datagen = ImageDataGenerator(
        featurewise_center=False,  # set input mean to 0 over the dataset
        samplewise_center=False,  # set each sample mean to 0
        featurewise_std_normalization=False,  # divide inputs by std of the dataset
        samplewise_std_normalization=False,  # divide each input by its std
        zca_whitening=False,  # apply ZCA whitening
        rotation_range=5,  # randomly rotate images in the range (degrees, 0 to 180)
        width_shift_range=0.1,  # randomly shift images horizontally (fraction of total width)
        height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
        horizontal_flip=False,  # randomly flip images
        vertical_flip=False)  # randomly flip images
    '''
    #datagen.fit(X_train)
    #Stop training when a monitored quantity has stopped improving.
    #monitor: quantity to be monitored.
    #patience: number of epochs with no improvement after which training will be stopped.
    #verbose: verbosity mode.
    #mode: one of {auto, min, max}. In 'min' mode, training will stop when the quantity monitored has stopped decreasing; in 'max' mode it will stop when the quantity monitored has stopped increasing.
    earlyStopping= EarlyStopping(monitor='val_loss', patience=5, verbose=1, mode='auto')
    checkpointer = ModelCheckpoint(filepath="../model/model/model_weights_"+str(i)+".h5", verbose=1, save_best_only=True)
#    class_weight = {0:1,1:5}
    #调用fit方法
    #shuffle=True，数据经过随机打乱
    #verbose=1，训练过程中输出的信息，1为输出进度条记录
    #validation_data，指定验证集
    model.fit(X_train, Y_train, batch_size=batch_size, nb_epoch=nb_epoch, shuffle=True, verbose=1, validation_data=(X_valid,Y_valid),callbacks=[checkpointer,earlyStopping])
    '''
    model.fit_generator(datagen.flow(X_train, Y_train,
                        batch_size=batch_size,
			shuffle=True),
                        samples_per_epoch=X_train.shape[0],
                        nb_epoch=nb_epoch,
			verbose=1,
                        validation_data=(X_valid, Y_valid),
			callbacks=[checkpointer,earlyStopping])
    '''
    #load the best model
    #model = model_from_json(open('../model/model/model_architecture_1.json').read())
    model.load_weights('../model/model/model_weights_'+str(i)+'.h5')

    score = model.evaluate(X_train, Y_train, verbose=1)
    print('Train loss:', score[0])
    print('Train accuracy:', score[1])

    count = 0
    result = model.predict(X_train, batch_size=100, verbose = 0)
    for j in range(X_train.shape[0]):
        if Y_train_label[j] == 1 and result[j][1] > 0.99:
            count += 1
        elif Y_train_label[j] == 0 and result[j][1] <= 0.99:
            count += 1
    print(('Strict Train accuracy:', float(count)/X_train.shape[0]))

    score = model.evaluate(X_valid, Y_valid, verbose=1)
    print('Validation loss:', score[0])
    print('Validation accuracy:', score[1])

    count = 0
    result = model.predict(X_valid, batch_size=100, verbose = 0)
    for j in range(X_valid.shape[0]):
        if Y_valid_label[j] == 1 and result[j][1] > 0.99:
            count += 1
        elif Y_valid_label[j] == 0 and result[j][1] <= 0.99:
            count += 1
    print(('Strict Validation accuracy:', float(count)/X_valid.shape[0]))

    '''
    score = model.evaluate(X_test, Y_test, verbose=1)
    print('Test loss:', score[0])
    print('Test accuracy:', score[1])

    count = 0
    result = model.predict(X_test, batch_size=100, verbose = 0)
    for j in range(X_test.shape[0]):
        if Y_test_label[j] == 1 and result[j][1] > 0.9:
            count += 1
        elif Y_test_label[j] == 0 and result[j][0] > 0.5:
            count += 1
    print(('Strict Test accuracy:', float(count)/X_test.shape[0]))
    '''
#    dist[i]=score[1]
    dist[i]=float(count)/X_valid.shape[0]

######################################
#保存CNN模型
######################################
path_save_model = '/home/cad/model'
json_string = model.to_json()
open(path_save_model+'/model_architecture.json','w').write(json_string)

##根据dist.values()做一个排序，选择前5个或10个效果较好的cnn model单独保存,应用于sliding window最后判断
print(dist)
import subprocess
for i in range(5):
        index = dist.keys()[dist.values().index(max(dist.values()))]
        subprocess.call(["mv", "../model/model/model_weights_"+str(index)+".h5",path_save_model+"/model_weights_"+str(index)+".h5"])
        dist.pop(index)
<commit_msg>fix import load_data<commit_after>#coding:utf-8
from __future__ import absolute_import
from __future__ import print_function
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Convolution2D, MaxPooling2D
from keras.optimizers import SGD
from keras.utils import np_utils, generic_utils
from six.moves import range
from keras.callbacks import EarlyStopping
from keras.callbacks import ModelCheckpoint
from keras.regularizers import l2, activity_l2
from keras.layers.advanced_activations import PReLU
from data import load_data
import h5py
import argparse
from keras.models import model_from_json
#import numpy as np
#np.random.seed(1337)  # for reproducibility
#from keras.utils.visualize_util import plot
#import tensorflow as tf
#import theano
#theano.config.device = 'gpu1'
#theano.config.floatX = 'float32'

#Number of samples per gradient update
batch_size = 100
#分类个数
nb_classes = 2
#the number of times to iterate over the training data arrays
nb_epoch = 20
#训练CNN的个数
nb_loop = 50

# input image dimensions
img_rows, img_cols = 64, 32
# the images are RGB
img_channels = 3
'''
ap = argparse.ArgumentParser()
ap.add_argument("--train", required=True, help="Path to the training set")
ap.add_argument("--valid", required=True, help="Path to the validation set")
ap.add_argument("--test", required=True, help="Path to the testing set")
ap.add_argument("--0", required=True, help="Path to the negative set")
args = vars(ap.parse_args())

#加载数据
path_train = args["train"]
path_valid = args["valid"]
path_test = args["test"]
path_0 = args["0"]
'''
#path_train='/home/cad/dataset/middlepink/data/train'
#path_valid='/home/cad/dataset/middlepink/data/valid_test_1'
#path_train_0='/home/cad/dataset/middlepink/train_negative'
#path_valid_0='/home/cad/dataset/middlepink/data/valid_test_0'

path_train='/home/cad/dataset/pinkall/train_1'
path_valid='/home/cad/dataset/pinkall/valid_test_1'
path_train_0='/home/cad/dataset/pinkall/train_0'
path_valid_0='/home/cad/dataset/pinkall/valid_test_0'

X_train,Y_train_label,X_valid,Y_valid_label = load_data(path_train,path_valid,path_train_0,path_valid_0)
print(X_train.shape[0], 'training samples')
print(X_valid.shape[0],'validation samples')
#print(X_test.shape[0],'testing samples')

#label为0,1两个类别，keras要求格式为binary class matrices,转化一下，调用keras提供的函数
Y_train = np_utils.to_categorical(Y_train_label, nb_classes)
Y_valid = np_utils.to_categorical(Y_valid_label, nb_classes)
#Y_test = np_utils.to_categorical(Y_test_label, nb_classes)

dist={}
#prelu=PReLU(init='zero', weights=None)
for i in range(nb_loop):
    ###############
    #开始建立CNN模型
    ###############
    model = Sequential()
    #第一个卷积层，32个卷积核，每个卷积核大小5*5
    #border_mode为valid
    #激活函数用relu
    model.add(Convolution2D(32, 5, 5,
            #        border_mode='valid',
            #        activation='relu',
            #        W_regularizer=l2(0.01),
            #        b_regularizer=l2(0.01),
            #       activity_regularizer=activity_l2(0.001),
                    input_shape=(img_channels, img_rows, img_cols)))
                    #input_shape=(X_train.shape[1:])))
    model.add(Activation('relu'))
#    model.add(prelu)
#    model.add(MaxPooling2D(pool_size=(2, 2)))
#    model.add(Dropout(0.5))

    #第二个卷积层，32个卷积核，每个卷积核大小5*5
    #采用maxpooling,poolsize为（2,2）
    #使用Dropout
    model.add(Convolution2D(32, 5, 5))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.5))

    #第三个卷积层，64个卷积核，每个卷积核大小3*3
    model.add(Convolution2D(64, 3, 3))
    model.add(Activation('relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.5))

    #全连接层
    model.add(Flatten())
    model.add(Dense(128))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))

    #Softmax回归，输出2个类别
    model.add(Dense(nb_classes))
    model.add(Activation('softmax'))

    ##############
    #开始训练模型
    ##############

    #lr: float >= 0. Learning rate.
    #momentum: float >= 0. Parameter updates momentum.
    #decay: float >= 0. Learning rate decay over each update.
    #nesterov: boolean. Whether to apply Nesterov momentum.
    #model.compile里的参数loss就是损失函数(目标函数),binary_crossentropy即为logistic loss
    sgd = SGD(lr=0.004, decay=1e-6, momentum=0.9, nesterov=True)
    model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])

    #this will do preprocessing and realtime data augmentation
    '''
    datagen = ImageDataGenerator(
        featurewise_center=False,  # set input mean to 0 over the dataset
        samplewise_center=False,  # set each sample mean to 0
        featurewise_std_normalization=False,  # divide inputs by std of the dataset
        samplewise_std_normalization=False,  # divide each input by its std
        zca_whitening=False,  # apply ZCA whitening
        rotation_range=5,  # randomly rotate images in the range (degrees, 0 to 180)
        width_shift_range=0.1,  # randomly shift images horizontally (fraction of total width)
        height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
        horizontal_flip=False,  # randomly flip images
        vertical_flip=False)  # randomly flip images
    '''
    #datagen.fit(X_train)
    #Stop training when a monitored quantity has stopped improving.
    #monitor: quantity to be monitored.
    #patience: number of epochs with no improvement after which training will be stopped.
    #verbose: verbosity mode.
    #mode: one of {auto, min, max}. In 'min' mode, training will stop when the quantity monitored has stopped decreasing; in 'max' mode it will stop when the quantity monitored has stopped increasing.
    earlyStopping= EarlyStopping(monitor='val_loss', patience=5, verbose=1, mode='auto')
    checkpointer = ModelCheckpoint(filepath="../model/model/model_weights_"+str(i)+".h5", verbose=1, save_best_only=True)
#    class_weight = {0:1,1:5}
    #调用fit方法
    #shuffle=True，数据经过随机打乱
    #verbose=1，训练过程中输出的信息，1为输出进度条记录
    #validation_data，指定验证集
    model.fit(X_train, Y_train, batch_size=batch_size, nb_epoch=nb_epoch, shuffle=True, verbose=1, validation_data=(X_valid,Y_valid),callbacks=[checkpointer,earlyStopping])
    '''
    model.fit_generator(datagen.flow(X_train, Y_train,
                        batch_size=batch_size,
			shuffle=True),
                        samples_per_epoch=X_train.shape[0],
                        nb_epoch=nb_epoch,
			verbose=1,
                        validation_data=(X_valid, Y_valid),
			callbacks=[checkpointer,earlyStopping])
    '''
    #load the best model
    #model = model_from_json(open('../model/model/model_architecture_1.json').read())
    model.load_weights('../model/model/model_weights_'+str(i)+'.h5')

    score = model.evaluate(X_train, Y_train, verbose=1)
    print('Train loss:', score[0])
    print('Train accuracy:', score[1])

    count = 0
    result = model.predict(X_train, batch_size=100, verbose = 0)
    for j in range(X_train.shape[0]):
        if Y_train_label[j] == 1 and result[j][1] > 0.99:
            count += 1
        elif Y_train_label[j] == 0 and result[j][1] <= 0.99:
            count += 1
    print(('Strict Train accuracy:', float(count)/X_train.shape[0]))

    score = model.evaluate(X_valid, Y_valid, verbose=1)
    print('Validation loss:', score[0])
    print('Validation accuracy:', score[1])

    count = 0
    result = model.predict(X_valid, batch_size=100, verbose = 0)
    for j in range(X_valid.shape[0]):
        if Y_valid_label[j] == 1 and result[j][1] > 0.99:
            count += 1
        elif Y_valid_label[j] == 0 and result[j][1] <= 0.99:
            count += 1
    print(('Strict Validation accuracy:', float(count)/X_valid.shape[0]))

    '''
    score = model.evaluate(X_test, Y_test, verbose=1)
    print('Test loss:', score[0])
    print('Test accuracy:', score[1])

    count = 0
    result = model.predict(X_test, batch_size=100, verbose = 0)
    for j in range(X_test.shape[0]):
        if Y_test_label[j] == 1 and result[j][1] > 0.9:
            count += 1
        elif Y_test_label[j] == 0 and result[j][0] > 0.5:
            count += 1
    print(('Strict Test accuracy:', float(count)/X_test.shape[0]))
    '''
#    dist[i]=score[1]
    dist[i]=float(count)/X_valid.shape[0]

######################################
#保存CNN模型
######################################
path_save_model = '/home/cad/model'
json_string = model.to_json()
open(path_save_model+'/model_architecture.json','w').write(json_string)

##根据dist.values()做一个排序，选择前5个或10个效果较好的cnn model单独保存,应用于sliding window最后判断
print(dist)
import subprocess
for i in range(5):
        index = dist.keys()[dist.values().index(max(dist.values()))]
        subprocess.call(["mv", "../model/model/model_weights_"+str(index)+".h5",path_save_model+"/model_weights_"+str(index)+".h5"])
        dist.pop(index)
<|endoftext|>


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

# file: $CUIBM_DIR/scripts/python/readData.py
# author: Krishnan, A. (anush@bu.edu)
# description: definitions of functions to read numerical data

import os
import sys
import argparse

import numpy as np
from numpy import linalg as la
from matplotlib import pyplt as plt

def readSimulationParameters(folder):
	"""Reads the parameters of the simulation.

	Arguments
	---------
	folder -- path of the simulation case.

	Returns
	-------
	nt -- total number of time-steps.
	start_step -- initial time-step.
	nsave -- saving interval.
	dt -- time-increment.
	"""
	# open and read the info file
	with open(folder+'/run.info', 'r') as infile:
		option_list = infile.read().split()

	# create the parser
	parser = argparse.ArgumentParser()
	# fill the parser with arguments
	parser.add_argument('--nx', dest='nx', type=int,
						help='number of cells in the x-direction')
	parser.add_argument('--ny', dest='ny', type=int,
						help='number of cells in the y-direction')
	parser.add_argument('--nt', dest='nt', type=int,
						help='total number of time-steps')
	parser.add_argument('--nsave', dest='nsave', type=int,
						help='saving interval')
	parser.add_argument('--startStep', dest='start_step', type=int,
						help='starting time-step')
	parser.add_argument('--dt', dest='dt', type=float,
						help='time-increment')
	parser.add_argument('--vortlim', dest='vort_lim', type=float,
						help='vorticity limit')
	parser.add_argument('--folder', dest='folder_name', type=str,
						help='name of the folder in which data is saved')
	parser.add_argument('--nu', dest='nu', type=float,
						help='kinematic viscosity')
	parser.add_argument('--input', dest='body_file', type=str,
						help='body input file')
	parser.add_argument('--flowFile', dest='flow_file', type=str,
						help='flow description file')
	parser.add_argument('--domainFile', dest='domain_file', type=str,
						help='mesh description file')
	parser.add_argument('--simulation_file', dest='simulation_file', type=str,
						help='simulation parameters file')
	parser.add_argument('--bodyFile', dest='body_file', type=str,
						help='body file')
	# parse the list
	args = parser.parse_args(option_list)

	return args.nt, args.start_step, args.nsave, args.dt

def readGridData(folder):
	# open the file with the grid data
	gridFile = folder + '/grid'
	f = open(gridFile, 'r')

	# read the number of cells in the x-direction
	a = f.readline().strip().split()
	nx = int(a[0])

	# arrays to store grid information
	x = np.zeros(nx+1)
	dx = np.zeros(nx)
	xu = np.zeros(nx-1)
	xv = np.zeros(nx)

	# read the x-coordinates of the grid nodes
	a = f.readline().strip().split()
	x[0] = float(a[0])
	for i in range(nx):
		a = f.readline().strip().split()
		if a==[]:
			break
		x[i+1] = float(a[0])
		dx[i] = x[i+1] - x[i]
		xv[i] = 0.5*(x[i+1] + x[i]);
		if i < nx-1:
			xu[i] = x[i+1]

	a = f.readline() # skip the empty lines

	# read the number of cells in the y-direction
	a = f.readline().strip().split()
	ny = int(a[0])

	# arrays to store y-grid information
	y = np.zeros(ny+1)
	dy = np.zeros(ny)
	yu = np.zeros(ny)
	yv = np.zeros(ny-1)

	# read the y-coordinates of the grid nodes
	a = f.readline().strip().split()
	y[0] = float(a[0])
	for j in range(ny):
		a = f.readline().strip().split()
		if a==[]:
			break
		y[j+1] = float(a[0])
		dy[j] = y[j+1] - y[j]
		yu[j] = 0.5*(y[j+1] + y[j])
		if j < ny-1:
			yv[j] = y[j+1]

	f.close()

	return nx, ny, dx, dy, xu, yu, xv, yv

def read_velocity_data(folder, time_step, nx, ny, dx, dy):
	u = np.empty((nx-1)*ny, dtype=float)
	v = np.empty(nx*(ny-1), dtype=float)

	flux_file = '%s/%07d/q' % (folder, time_step)
	with open(flux_file, 'r') as infile:

def readVelocityData(folder, timeStep, nx, ny, dx, dy):
	u = np.zeros((nx-1)*ny)  # vector to store U
	v = np.zeros(nx*(ny-1))  # vector to store V

	fluxFile = folder + '/' + ("%07d" % timeStep) + "/q"
	try:
		f = open(fluxFile, 'r')
	except IOError:
		print "File %s missing!" % fluxFile
		return None, None

	# length of vector q
	a = f.readline().strip().split()
	nq = int(a[0])

	# read the u velocities
	for j in range(ny):
		for i in range(nx-1):
			a = f.readline().strip().split()
			u[j*(nx-1)+i] = float(a[0])/dy[j]

	# read the u velocities
	for j in range(ny-1):
		for i in range(nx):
			a = f.readline().strip().split()
			v[j*nx+i] = float(a[0])/dx[i]

	return u, v
<commit_msg>Rewrite function readVelocityData()<commit_after>#!/usr/bin/env python

# file: $CUIBM_DIR/scripts/python/readData.py
# author: Krishnan, A. (anush@bu.edu)
# description: definitions of functions to read numerical data

import argparse

import numpy as np

def readSimulationParameters(folder):
	"""Reads the parameters of the simulation.

	Arguments
	---------
	folder -- path of the simulation case.

	Returns
	-------
	nt -- total number of time-steps.
	start_step -- initial time-step.
	nsave -- saving interval.
	dt -- time-increment.
	"""
	# open and read the info file
	with open(folder+'/run.info', 'r') as infile:
		option_list = infile.read().split()

	# create the parser
	parser = argparse.ArgumentParser()
	# fill the parser with arguments
	parser.add_argument('--nx', dest='nx', type=int,
						help='number of cells in the x-direction')
	parser.add_argument('--ny', dest='ny', type=int,
						help='number of cells in the y-direction')
	parser.add_argument('--nt', dest='nt', type=int,
						help='total number of time-steps')
	parser.add_argument('--nsave', dest='nsave', type=int,
						help='saving interval')
	parser.add_argument('--startStep', dest='start_step', type=int,
						help='starting time-step')
	parser.add_argument('--dt', dest='dt', type=float,
						help='time-increment')
	parser.add_argument('--vortlim', dest='vort_lim', type=float,
						help='vorticity limit')
	parser.add_argument('--folder', dest='folder_name', type=str,
						help='name of the folder in which data is saved')
	parser.add_argument('--nu', dest='nu', type=float,
						help='kinematic viscosity')
	parser.add_argument('--input', dest='body_file', type=str,
						help='body input file')
	parser.add_argument('--flowFile', dest='flow_file', type=str,
						help='flow description file')
	parser.add_argument('--domainFile', dest='domain_file', type=str,
						help='mesh description file')
	parser.add_argument('--simulation_file', dest='simulation_file', type=str,
						help='simulation parameters file')
	parser.add_argument('--bodyFile', dest='body_file', type=str,
						help='body file')
	# parse the list
	args = parser.parse_args(option_list)

	return args.nt, args.start_step, args.nsave, args.dt

def readGridData(folder):
	# open the file with the grid data
	gridFile = folder + '/grid'
	f = open(gridFile, 'r')

	# read the number of cells in the x-direction
	a = f.readline().strip().split()
	nx = int(a[0])

	# arrays to store grid information
	x = np.zeros(nx+1)
	dx = np.zeros(nx)
	xu = np.zeros(nx-1)
	xv = np.zeros(nx)

	# read the x-coordinates of the grid nodes
	a = f.readline().strip().split()
	x[0] = float(a[0])
	for i in range(nx):
		a = f.readline().strip().split()
		if a==[]:
			break
		x[i+1] = float(a[0])
		dx[i] = x[i+1] - x[i]
		xv[i] = 0.5*(x[i+1] + x[i]);
		if i < nx-1:
			xu[i] = x[i+1]

	a = f.readline() # skip the empty lines

	# read the number of cells in the y-direction
	a = f.readline().strip().split()
	ny = int(a[0])

	# arrays to store y-grid information
	y = np.zeros(ny+1)
	dy = np.zeros(ny)
	yu = np.zeros(ny)
	yv = np.zeros(ny-1)

	# read the y-coordinates of the grid nodes
	a = f.readline().strip().split()
	y[0] = float(a[0])
	for j in range(ny):
		a = f.readline().strip().split()
		if a==[]:
			break
		y[j+1] = float(a[0])
		dy[j] = y[j+1] - y[j]
		yu[j] = 0.5*(y[j+1] + y[j])
		if j < ny-1:
			yv[j] = y[j+1]

	f.close()

	return nx, ny, dx, dy, xu, yu, xv, yv

def readVelocityData(folder, time_step, nx, ny, dx, dy):
	"""Reads the velocity data at a given time-step.

	Arguments
	---------
	folder -- path of the simulation case.
	time-step -- current time-step.
	nx, ny -- number of cells on the x- and y- directions.
	dx, dy -- cell widths in the x- and y- directions.
	"""
	flux_file = '%s/%07d/q' % (folder, time_step)
	with open(flux_file, 'r') as infile:
		nq = int(infile.readline())			# length of flux vector q
		q = np.loadtxt(infile, dtype=float)	# store flux vector q

	# store u-velocities
	n_u = (nx-1) * ny				# number of u-velocity points
	u = np.empty(n_u, dtype=float)
	for j in xrange(ny):
		for i in xrange(nx-1):
			u[j*(nx-1)+i] = q[j*(nx-1)+i]/dy[j]

	# store v-velocities
	v = np.empty(nx*(ny-1), dtype=float)
	for j in xrange(ny-1):
		for i in xrange(nx):
			v[j*nx+i] = q[n_u+j*nx+i]/dx[i]

	return u, v
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Event Dispatcher

Implements an priority queue based event dispatcher.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import collections
import heapq
import itertools

from twisted.internet import defer
from twisted.logger import Logger

Key = collections.namedtuple('Key', ['priority', 'serial'])
Entry = collections.namedtuple('Entry', ['key', 'callback'])

class EventDispatcher(object):
    log = Logger()

    def __init__(self):
        self._active = set()
        self._counter = itertools.count()
        self._listeners = {}

    def add_listener(self, event_type, priority, callback):
        listeners = self._listeners.setdefault(event_type, [])
        key = Key(priority, next(self._counter))
        heapq.heappush(listeners, Entry(key, callback))
        self._active.add(key)
        return key

    def remove_listener(self, key):
        self._active.remove(key)

    @defer.inlineCallbacks
    def dispatch(self, event, logfails=False):
        event_type = type(event)
        if event_type in self._listeners:
            listeners = self._listeners[event_type]
            keyfunc = lambda entry: entry.key.priority

            for priority, group in itertools.groupby(listeners, keyfunc):
                batch = []

                self.log.debug('Calling handlers with priority {priority} for {event}', event=event, priority=priority)
                for key, callback in group:
                    if key in self._active:
                        d = defer.maybeDeferred(callback, event)
                        if logfails:
                            d.addErrback(self._logfail, 'Failure ignored while handling event {event}', event=event)
                        batch.append(d)

                if len(batch):
                    yield defer.DeferredList(batch, fireOnOneErrback=True)

                self.log.debug('Called {n} handlers with priority {priority} for {event}', n=len(batch), event=event, priority=priority)

    def prune(self):
        pruned_listeners = {}
        for event_type, listeners in self._listeners.items():
            listeners = [(key, callback) for key, callback in listeners if key in self._active]
            if len(listeners):
                pruned_listeners[event_type] = listeners

        self._listeners = pruned_listeners

    def _logfail(self, failure, fmt, *args, **kwds):
        """
        Errback: Logs and consumes a failure.
        """
        self.log.failure(fmt, failure, *args, **kwds)
<commit_msg>Keep sequence of event handlers, pass *args and **kwds when invoking events, extract get_listeners method<commit_after># -*- coding: utf-8 -*-
"""Event Dispatcher

Implements an priority queue based event dispatcher.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import collections
import itertools
import bisect

from twisted.internet import defer
from twisted.logger import Logger

Key = collections.namedtuple('Key', ['priority', 'serial'])
Handler = collections.namedtuple('Handler', ['callback', 'args', 'kwds'])
Entry = collections.namedtuple('Entry', ['key', 'handler'])

class EventDispatcher(object):
    log = Logger()

    def __init__(self):
        self._active = set()
        self._counter = itertools.count()
        self._listeners = {}

    def add_listener(self, event_type, priority, callback, *args, **kwds):
        listeners = self._listeners.setdefault(event_type, [])
        key = Key(priority, next(self._counter))
        bisect.insort(listeners, Entry(key, Handler(callback, args, kwds)))
        self._active.add(key)
        return key

    def remove_listener(self, key):
        self._active.remove(key)

    def get_listeners(self, event_type):
        if event_type in self._listeners:
            keyfunc = lambda entry: entry.key.priority
            filterfunc = lambda entry: entry.key in self._active

            event_listeners = self._listeners[event_type]
            filtered_listeners = itertools.ifilter(filterfunc, event_listeners)
            return itertools.groupby(filtered_listeners, keyfunc)
        else:
            return iter([])

    @defer.inlineCallbacks
    def dispatch(self, event, logfails=False):
        for priority, group in self.get_listeners(type(event)):
            self.log.debug('Calling handlers with priority {priority} for {event}', event=event, priority=priority)

            batch = []
            for key, handler in group:
                d = defer.maybeDeferred(handler.callback, event, *handler.args, **handler.kwds)
                if logfails:
                    d.addErrback(self._logfail, 'Failure ignored while handling event {event} with {handler}', event=event, priority=priority, key=key, handler=handler)
                batch.append(d)

            if len(batch):
                yield defer.DeferredList(batch, fireOnOneErrback=True)

            self.log.debug('Called {n} handlers with priority {priority} for {event}', n=len(batch), event=event, priority=priority)

    def prune(self):
        pruned_listeners = {}
        for event_type, listeners in self._listeners.items():
            listeners = [(key, handler) for key, handler in listeners if key in self._active]
            if len(listeners):
                pruned_listeners[event_type] = listeners

        self._listeners = pruned_listeners

    def _logfail(self, failure, fmt, *args, **kwds):
        """
        Errback: Logs and consumes a failure.
        """
        self.log.failure(fmt, failure, *args, **kwds)
<|endoftext|>


<commit_before>'''SSL with SNI_-support for Python 2. Follow these instructions if you would
like to verify SSL certificates in Python 2. Note, the default libraries do
*not* do certificate checking; you need to do additional work to validate
certificates yourself.

This needs the following packages installed:

* pyOpenSSL (tested with 0.13)
* ndg-httpsclient (tested with 0.3.2)
* pyasn1 (tested with 0.1.6)

You can install them with the following command:

    pip install pyopenssl ndg-httpsclient pyasn1

To activate certificate checking, call
:func:`~urllib3.contrib.pyopenssl.inject_into_urllib3` from your Python code
before you begin making HTTP requests. This can be done in a ``sitecustomize``
module, or at any other time before your application begins using ``urllib3``,
like this::

    try:
        import urllib3.contrib.pyopenssl
        urllib3.contrib.pyopenssl.inject_into_urllib3()
    except ImportError:
        pass

Now you can use :mod:`urllib3` as you normally would, and it will support SNI
when the required modules are installed.

Activating this module also has the positive side effect of disabling SSL/TLS
compression in Python 2 (see `CRIME attack`_).

If you want to configure the default list of supported cipher suites, you can
set the ``urllib3.contrib.pyopenssl.DEFAULT_SSL_CIPHER_LIST`` variable.

Module Variables
----------------

:var DEFAULT_SSL_CIPHER_LIST: The list of supported SSL/TLS cipher suites.

.. _sni: https://en.wikipedia.org/wiki/Server_Name_Indication
.. _crime attack: https://en.wikipedia.org/wiki/CRIME_(security_exploit)

'''
from __future__ import absolute_import

try:
    from ndg.httpsclient.ssl_peer_verification import SUBJ_ALT_NAME_SUPPORT
    from ndg.httpsclient.subj_alt_name import SubjectAltName as BaseSubjectAltName
except SyntaxError as e:
    raise ImportError(e)

import OpenSSL.SSL
from pyasn1.codec.der import decoder as der_decoder
from pyasn1.type import univ, constraint
from socket import timeout, error as SocketError

try:  # Platform-specific: Python 2
    from socket import _fileobject
except ImportError:  # Platform-specific: Python 3
    _fileobject = None
    from urllib3.packages.backports.makefile import backport_makefile

import ssl
import select
import six

from .. import connection
from .. import util

__all__ = ['inject_into_urllib3', 'extract_from_urllib3']

# SNI only *really* works if we can read the subjectAltName of certificates.
HAS_SNI = SUBJ_ALT_NAME_SUPPORT

# Map from urllib3 to PyOpenSSL compatible parameter-values.
_openssl_versions = {
    ssl.PROTOCOL_SSLv23: OpenSSL.SSL.SSLv23_METHOD,
    ssl.PROTOCOL_TLSv1: OpenSSL.SSL.TLSv1_METHOD,
}

if hasattr(ssl, 'PROTOCOL_TLSv1_1') and hasattr(OpenSSL.SSL, 'TLSv1_1_METHOD'):
    _openssl_versions[ssl.PROTOCOL_TLSv1_1] = OpenSSL.SSL.TLSv1_1_METHOD

if hasattr(ssl, 'PROTOCOL_TLSv1_2') and hasattr(OpenSSL.SSL, 'TLSv1_2_METHOD'):
    _openssl_versions[ssl.PROTOCOL_TLSv1_2] = OpenSSL.SSL.TLSv1_2_METHOD

try:
    _openssl_versions.update({ssl.PROTOCOL_SSLv3: OpenSSL.SSL.SSLv3_METHOD})
except AttributeError:
    pass

_openssl_verify = {
    ssl.CERT_NONE: OpenSSL.SSL.VERIFY_NONE,
    ssl.CERT_OPTIONAL: OpenSSL.SSL.VERIFY_PEER,
    ssl.CERT_REQUIRED:
        OpenSSL.SSL.VERIFY_PEER + OpenSSL.SSL.VERIFY_FAIL_IF_NO_PEER_CERT,
}

DEFAULT_SSL_CIPHER_LIST = util.ssl_.DEFAULT_CIPHERS

# OpenSSL will only write 16K at a time
SSL_WRITE_BLOCKSIZE = 16384

orig_util_HAS_SNI = util.HAS_SNI
orig_connection_ssl_wrap_socket = connection.ssl_wrap_socket

def inject_into_urllib3():
    'Monkey-patch urllib3 with PyOpenSSL-backed SSL-support.'

    connection.ssl_wrap_socket = ssl_wrap_socket
    util.HAS_SNI = HAS_SNI
    util.IS_PYOPENSSL = True

def extract_from_urllib3():
    'Undo monkey-patching by :func:`inject_into_urllib3`.'

    connection.ssl_wrap_socket = orig_connection_ssl_wrap_socket
    util.HAS_SNI = orig_util_HAS_SNI
    util.IS_PYOPENSSL = False

# Note: This is a slightly bug-fixed version of same from ndg-httpsclient.
class SubjectAltName(BaseSubjectAltName):
    '''ASN.1 implementation for subjectAltNames support'''

    # There is no limit to how many SAN certificates a certificate may have,
    #   however this needs to have some limit so we'll set an arbitrarily high
    #   limit.
    sizeSpec = univ.SequenceOf.sizeSpec + \
        constraint.ValueSizeConstraint(1, 1024)

# Note: This is a slightly bug-fixed version of same from ndg-httpsclient.
def get_subj_alt_name(peer_cert):
    # Search through extensions
    dns_name = []
    if not SUBJ_ALT_NAME_SUPPORT:
        return dns_name

    general_names = SubjectAltName()
    for i in range(peer_cert.get_extension_count()):
        ext = peer_cert.get_extension(i)
        ext_name = ext.get_short_name()
        if ext_name != b'subjectAltName':
            continue

        # PyOpenSSL returns extension data in ASN.1 encoded form
        ext_dat = ext.get_data()
        decoded_dat = der_decoder.decode(ext_dat,
                                         asn1Spec=general_names)

        for name in decoded_dat:
            if not isinstance(name, SubjectAltName):
                continue
            for entry in range(len(name)):
                component = name.getComponentByPosition(entry)
                if component.getName() != 'dNSName':
                    continue
                dns_name.append(str(component.getComponent()))

    return dns_name

class WrappedSocket(object):
    '''API-compatibility wrapper for Python OpenSSL's Connection-class.

    Note: _makefile_refs, _drop() and _reuse() are needed for the garbage
    collector of pypy.
    '''

    def __init__(self, connection, socket, suppress_ragged_eofs=True):
        self.connection = connection
        self.socket = socket
        self.suppress_ragged_eofs = suppress_ragged_eofs
        self._makefile_refs = 0
        self._closed = False

    def fileno(self):
        return self.socket.fileno()

    if _fileobject:  # Platform-specific: Python 2
        def makefile(self, mode, bufsize=-1):
            self._makefile_refs += 1
            return _fileobject(self, mode, bufsize, close=True)
    else:  # Platform-specific: Python 3
        WrappedSocket.makefile = backport_makefile

        # Copy-pasted from Python 3.5 source code
        def _decref_socketios(self):
            if self._makefile_refs > 0:
                self._makefile_refs -= 1
            if self._closed:
                self.close()

    def recv(self, *args, **kwargs):
        try:
            data = self.connection.recv(*args, **kwargs)
        except OpenSSL.SSL.SysCallError as e:
            if self.suppress_ragged_eofs and e.args == (-1, 'Unexpected EOF'):
                return b''
            else:
                raise SocketError(e)
        except OpenSSL.SSL.ZeroReturnError as e:
            if self.connection.get_shutdown() == OpenSSL.SSL.RECEIVED_SHUTDOWN:
                return b''
            else:
                raise
        except OpenSSL.SSL.WantReadError:
            rd, wd, ed = select.select(
                [self.socket], [], [], self.socket.gettimeout())
            if not rd:
                raise timeout('The read operation timed out')
            else:
                return self.recv(*args, **kwargs)
        else:
            return data

    def recv_into(self, *args, **kwargs):
        try:
            return self.connection.recv_into(*args, **kwargs)
        except OpenSSL.SSL.SysCallError as e:
            if self.suppress_ragged_eofs and e.args == (-1, 'Unexpected EOF'):
                return 0
            else:
                raise SocketError(str(e))
        except OpenSSL.SSL.ZeroReturnError as e:
            if self.connection.get_shutdown() == OpenSSL.SSL.RECEIVED_SHUTDOWN:
                return 0
            else:
                raise
        except OpenSSL.SSL.WantReadError:
            rd, wd, ed = select.select(
                [self.socket], [], [], self.socket.gettimeout())
            if not rd:
                raise timeout('The read operation timed out')
            else:
                return self.recv_into(*args, **kwargs)

    def settimeout(self, timeout):
        return self.socket.settimeout(timeout)

    def _send_until_done(self, data):
        while True:
            try:
                return self.connection.send(data)
            except OpenSSL.SSL.WantWriteError:
                _, wlist, _ = select.select([], [self.socket], [],
                                            self.socket.gettimeout())
                if not wlist:
                    raise timeout()
                continue

    def sendall(self, data):
        total_sent = 0
        while total_sent < len(data):
            sent = self._send_until_done(data[total_sent:total_sent + SSL_WRITE_BLOCKSIZE])
            total_sent += sent

    def shutdown(self):
        # FIXME rethrow compatible exceptions should we ever use this
        self.connection.shutdown()

    def close(self):
        if self._makefile_refs < 1:
            try:
                self._closed = True
                return self.connection.close()
            except OpenSSL.SSL.Error:
                return
        else:
            self._makefile_refs -= 1

    def getpeercert(self, binary_form=False):
        x509 = self.connection.get_peer_certificate()

        if not x509:
            return x509

        if binary_form:
            return OpenSSL.crypto.dump_certificate(
                OpenSSL.crypto.FILETYPE_ASN1,
                x509)

        return {
            'subject': (
                (('commonName', x509.get_subject().CN),),
            ),
            'subjectAltName': [
                ('DNS', value)
                for value in get_subj_alt_name(x509)
            ]
        }

    def _reuse(self):
        self._makefile_refs += 1

    def _drop(self):
        if self._makefile_refs < 1:
            self.close()
        else:
            self._makefile_refs -= 1

def _verify_callback(cnx, x509, err_no, err_depth, return_code):
    return err_no == 0

def ssl_wrap_socket(sock, keyfile=None, certfile=None, cert_reqs=None,
                    ca_certs=None, server_hostname=None,
                    ssl_version=None, ca_cert_dir=None):
    ctx = OpenSSL.SSL.Context(_openssl_versions[ssl_version])
    if certfile:
        keyfile = keyfile or certfile  # Match behaviour of the normal python ssl library
        ctx.use_certificate_file(certfile)
    if keyfile:
        ctx.use_privatekey_file(keyfile)
    if cert_reqs != ssl.CERT_NONE:
        ctx.set_verify(_openssl_verify[cert_reqs], _verify_callback)
    if ca_certs or ca_cert_dir:
        try:
            ctx.load_verify_locations(ca_certs, ca_cert_dir)
        except OpenSSL.SSL.Error as e:
            raise ssl.SSLError('bad ca_certs: %r' % ca_certs, e)
    else:
        ctx.set_default_verify_paths()

    # Disable TLS compression to mitigate CRIME attack (issue #309)
    OP_NO_COMPRESSION = 0x20000
    ctx.set_options(OP_NO_COMPRESSION)

    # Set list of supported ciphersuites.
    ctx.set_cipher_list(DEFAULT_SSL_CIPHER_LIST)

    cnx = OpenSSL.SSL.Connection(ctx, sock)
    if isinstance(server_hostname, six.text_type):  # Platform-specific: Python 3
        server_hostname = server_hostname.encode('utf-8')
    cnx.set_tlsext_host_name(server_hostname)
    cnx.set_connect_state()
    while True:
        try:
            cnx.do_handshake()
        except OpenSSL.SSL.WantReadError:
            rd, _, _ = select.select([sock], [], [], sock.gettimeout())
            if not rd:
                raise timeout('select timed out')
            continue
        except OpenSSL.SSL.Error as e:
            raise ssl.SSLError('bad handshake: %r' % e)
        break

    return WrappedSocket(cnx, sock)
<commit_msg>Scoping is tricky.<commit_after>'''SSL with SNI_-support for Python 2. Follow these instructions if you would
like to verify SSL certificates in Python 2. Note, the default libraries do
*not* do certificate checking; you need to do additional work to validate
certificates yourself.

This needs the following packages installed:

* pyOpenSSL (tested with 0.13)
* ndg-httpsclient (tested with 0.3.2)
* pyasn1 (tested with 0.1.6)

You can install them with the following command:

    pip install pyopenssl ndg-httpsclient pyasn1

To activate certificate checking, call
:func:`~urllib3.contrib.pyopenssl.inject_into_urllib3` from your Python code
before you begin making HTTP requests. This can be done in a ``sitecustomize``
module, or at any other time before your application begins using ``urllib3``,
like this::

    try:
        import urllib3.contrib.pyopenssl
        urllib3.contrib.pyopenssl.inject_into_urllib3()
    except ImportError:
        pass

Now you can use :mod:`urllib3` as you normally would, and it will support SNI
when the required modules are installed.

Activating this module also has the positive side effect of disabling SSL/TLS
compression in Python 2 (see `CRIME attack`_).

If you want to configure the default list of supported cipher suites, you can
set the ``urllib3.contrib.pyopenssl.DEFAULT_SSL_CIPHER_LIST`` variable.

Module Variables
----------------

:var DEFAULT_SSL_CIPHER_LIST: The list of supported SSL/TLS cipher suites.

.. _sni: https://en.wikipedia.org/wiki/Server_Name_Indication
.. _crime attack: https://en.wikipedia.org/wiki/CRIME_(security_exploit)

'''
from __future__ import absolute_import

try:
    from ndg.httpsclient.ssl_peer_verification import SUBJ_ALT_NAME_SUPPORT
    from ndg.httpsclient.subj_alt_name import SubjectAltName as BaseSubjectAltName
except SyntaxError as e:
    raise ImportError(e)

import OpenSSL.SSL
from pyasn1.codec.der import decoder as der_decoder
from pyasn1.type import univ, constraint
from socket import timeout, error as SocketError

try:  # Platform-specific: Python 2
    from socket import _fileobject
except ImportError:  # Platform-specific: Python 3
    _fileobject = None
    from urllib3.packages.backports.makefile import backport_makefile

import ssl
import select
import six

from .. import connection
from .. import util

__all__ = ['inject_into_urllib3', 'extract_from_urllib3']

# SNI only *really* works if we can read the subjectAltName of certificates.
HAS_SNI = SUBJ_ALT_NAME_SUPPORT

# Map from urllib3 to PyOpenSSL compatible parameter-values.
_openssl_versions = {
    ssl.PROTOCOL_SSLv23: OpenSSL.SSL.SSLv23_METHOD,
    ssl.PROTOCOL_TLSv1: OpenSSL.SSL.TLSv1_METHOD,
}

if hasattr(ssl, 'PROTOCOL_TLSv1_1') and hasattr(OpenSSL.SSL, 'TLSv1_1_METHOD'):
    _openssl_versions[ssl.PROTOCOL_TLSv1_1] = OpenSSL.SSL.TLSv1_1_METHOD

if hasattr(ssl, 'PROTOCOL_TLSv1_2') and hasattr(OpenSSL.SSL, 'TLSv1_2_METHOD'):
    _openssl_versions[ssl.PROTOCOL_TLSv1_2] = OpenSSL.SSL.TLSv1_2_METHOD

try:
    _openssl_versions.update({ssl.PROTOCOL_SSLv3: OpenSSL.SSL.SSLv3_METHOD})
except AttributeError:
    pass

_openssl_verify = {
    ssl.CERT_NONE: OpenSSL.SSL.VERIFY_NONE,
    ssl.CERT_OPTIONAL: OpenSSL.SSL.VERIFY_PEER,
    ssl.CERT_REQUIRED:
        OpenSSL.SSL.VERIFY_PEER + OpenSSL.SSL.VERIFY_FAIL_IF_NO_PEER_CERT,
}

DEFAULT_SSL_CIPHER_LIST = util.ssl_.DEFAULT_CIPHERS

# OpenSSL will only write 16K at a time
SSL_WRITE_BLOCKSIZE = 16384

orig_util_HAS_SNI = util.HAS_SNI
orig_connection_ssl_wrap_socket = connection.ssl_wrap_socket

def inject_into_urllib3():
    'Monkey-patch urllib3 with PyOpenSSL-backed SSL-support.'

    connection.ssl_wrap_socket = ssl_wrap_socket
    util.HAS_SNI = HAS_SNI
    util.IS_PYOPENSSL = True

def extract_from_urllib3():
    'Undo monkey-patching by :func:`inject_into_urllib3`.'

    connection.ssl_wrap_socket = orig_connection_ssl_wrap_socket
    util.HAS_SNI = orig_util_HAS_SNI
    util.IS_PYOPENSSL = False

# Note: This is a slightly bug-fixed version of same from ndg-httpsclient.
class SubjectAltName(BaseSubjectAltName):
    '''ASN.1 implementation for subjectAltNames support'''

    # There is no limit to how many SAN certificates a certificate may have,
    #   however this needs to have some limit so we'll set an arbitrarily high
    #   limit.
    sizeSpec = univ.SequenceOf.sizeSpec + \
        constraint.ValueSizeConstraint(1, 1024)

# Note: This is a slightly bug-fixed version of same from ndg-httpsclient.
def get_subj_alt_name(peer_cert):
    # Search through extensions
    dns_name = []
    if not SUBJ_ALT_NAME_SUPPORT:
        return dns_name

    general_names = SubjectAltName()
    for i in range(peer_cert.get_extension_count()):
        ext = peer_cert.get_extension(i)
        ext_name = ext.get_short_name()
        if ext_name != b'subjectAltName':
            continue

        # PyOpenSSL returns extension data in ASN.1 encoded form
        ext_dat = ext.get_data()
        decoded_dat = der_decoder.decode(ext_dat,
                                         asn1Spec=general_names)

        for name in decoded_dat:
            if not isinstance(name, SubjectAltName):
                continue
            for entry in range(len(name)):
                component = name.getComponentByPosition(entry)
                if component.getName() != 'dNSName':
                    continue
                dns_name.append(str(component.getComponent()))

    return dns_name

class WrappedSocket(object):
    '''API-compatibility wrapper for Python OpenSSL's Connection-class.

    Note: _makefile_refs, _drop() and _reuse() are needed for the garbage
    collector of pypy.
    '''

    def __init__(self, connection, socket, suppress_ragged_eofs=True):
        self.connection = connection
        self.socket = socket
        self.suppress_ragged_eofs = suppress_ragged_eofs
        self._makefile_refs = 0
        self._closed = False

    def fileno(self):
        return self.socket.fileno()

    # Copy-pasted from Python 3.5 source code
    def _decref_socketios(self):
        if self._makefile_refs > 0:
            self._makefile_refs -= 1
        if self._closed:
            self.close()

    def recv(self, *args, **kwargs):
        try:
            data = self.connection.recv(*args, **kwargs)
        except OpenSSL.SSL.SysCallError as e:
            if self.suppress_ragged_eofs and e.args == (-1, 'Unexpected EOF'):
                return b''
            else:
                raise SocketError(e)
        except OpenSSL.SSL.ZeroReturnError as e:
            if self.connection.get_shutdown() == OpenSSL.SSL.RECEIVED_SHUTDOWN:
                return b''
            else:
                raise
        except OpenSSL.SSL.WantReadError:
            rd, wd, ed = select.select(
                [self.socket], [], [], self.socket.gettimeout())
            if not rd:
                raise timeout('The read operation timed out')
            else:
                return self.recv(*args, **kwargs)
        else:
            return data

    def recv_into(self, *args, **kwargs):
        try:
            return self.connection.recv_into(*args, **kwargs)
        except OpenSSL.SSL.SysCallError as e:
            if self.suppress_ragged_eofs and e.args == (-1, 'Unexpected EOF'):
                return 0
            else:
                raise SocketError(str(e))
        except OpenSSL.SSL.ZeroReturnError as e:
            if self.connection.get_shutdown() == OpenSSL.SSL.RECEIVED_SHUTDOWN:
                return 0
            else:
                raise
        except OpenSSL.SSL.WantReadError:
            rd, wd, ed = select.select(
                [self.socket], [], [], self.socket.gettimeout())
            if not rd:
                raise timeout('The read operation timed out')
            else:
                return self.recv_into(*args, **kwargs)

    def settimeout(self, timeout):
        return self.socket.settimeout(timeout)

    def _send_until_done(self, data):
        while True:
            try:
                return self.connection.send(data)
            except OpenSSL.SSL.WantWriteError:
                _, wlist, _ = select.select([], [self.socket], [],
                                            self.socket.gettimeout())
                if not wlist:
                    raise timeout()
                continue

    def sendall(self, data):
        total_sent = 0
        while total_sent < len(data):
            sent = self._send_until_done(data[total_sent:total_sent + SSL_WRITE_BLOCKSIZE])
            total_sent += sent

    def shutdown(self):
        # FIXME rethrow compatible exceptions should we ever use this
        self.connection.shutdown()

    def close(self):
        if self._makefile_refs < 1:
            try:
                self._closed = True
                return self.connection.close()
            except OpenSSL.SSL.Error:
                return
        else:
            self._makefile_refs -= 1

    def getpeercert(self, binary_form=False):
        x509 = self.connection.get_peer_certificate()

        if not x509:
            return x509

        if binary_form:
            return OpenSSL.crypto.dump_certificate(
                OpenSSL.crypto.FILETYPE_ASN1,
                x509)

        return {
            'subject': (
                (('commonName', x509.get_subject().CN),),
            ),
            'subjectAltName': [
                ('DNS', value)
                for value in get_subj_alt_name(x509)
            ]
        }

    def _reuse(self):
        self._makefile_refs += 1

    def _drop(self):
        if self._makefile_refs < 1:
            self.close()
        else:
            self._makefile_refs -= 1

if _fileobject:  # Platform-specific: Python 2
    def makefile(self, mode, bufsize=-1):
        self._makefile_refs += 1
        return _fileobject(self, mode, bufsize, close=True)
else:  # Platform-specific: Python 3
    makefile = backport_makefile

WrappedSocket.makefile = makefile

def _verify_callback(cnx, x509, err_no, err_depth, return_code):
    return err_no == 0

def ssl_wrap_socket(sock, keyfile=None, certfile=None, cert_reqs=None,
                    ca_certs=None, server_hostname=None,
                    ssl_version=None, ca_cert_dir=None):
    ctx = OpenSSL.SSL.Context(_openssl_versions[ssl_version])
    if certfile:
        keyfile = keyfile or certfile  # Match behaviour of the normal python ssl library
        ctx.use_certificate_file(certfile)
    if keyfile:
        ctx.use_privatekey_file(keyfile)
    if cert_reqs != ssl.CERT_NONE:
        ctx.set_verify(_openssl_verify[cert_reqs], _verify_callback)
    if ca_certs or ca_cert_dir:
        try:
            ctx.load_verify_locations(ca_certs, ca_cert_dir)
        except OpenSSL.SSL.Error as e:
            raise ssl.SSLError('bad ca_certs: %r' % ca_certs, e)
    else:
        ctx.set_default_verify_paths()

    # Disable TLS compression to mitigate CRIME attack (issue #309)
    OP_NO_COMPRESSION = 0x20000
    ctx.set_options(OP_NO_COMPRESSION)

    # Set list of supported ciphersuites.
    ctx.set_cipher_list(DEFAULT_SSL_CIPHER_LIST)

    cnx = OpenSSL.SSL.Connection(ctx, sock)
    if isinstance(server_hostname, six.text_type):  # Platform-specific: Python 3
        server_hostname = server_hostname.encode('utf-8')
    cnx.set_tlsext_host_name(server_hostname)
    cnx.set_connect_state()
    while True:
        try:
            cnx.do_handshake()
        except OpenSSL.SSL.WantReadError:
            rd, _, _ = select.select([sock], [], [], sock.gettimeout())
            if not rd:
                raise timeout('select timed out')
            continue
        except OpenSSL.SSL.Error as e:
            raise ssl.SSLError('bad handshake: %r' % e)
        break

    return WrappedSocket(cnx, sock)
<|endoftext|>