<commit_before># Copyright (C) 2007-2011 Michael Foord & the mock team
# E-mail: fuzzyman AT voidspace DOT org DOT uk
# http://www.voidspace.org.uk/python/mock/

from __future__ import with_statement

from tests.support import unittest2, is_instance

from mock import MagicMock, Mock, patch, sentinel

from tests.support_with import catch_warnings, nested

something  = sentinel.Something
something_else  = sentinel.SomethingElse

class WithTest(unittest2.TestCase):

    def testWithStatement(self):
        with patch('tests._testwith.something', sentinel.Something2):
            self.assertEqual(something, sentinel.Something2, "unpatched")
        self.assertEqual(something, sentinel.Something)

    def testWithStatementException(self):
        try:
            with patch('tests._testwith.something', sentinel.Something2):
                self.assertEqual(something, sentinel.Something2, "unpatched")
                raise Exception('pow')
        except Exception:
            pass
        else:
            self.fail("patch swallowed exception")
        self.assertEqual(something, sentinel.Something)

    def testWithStatementAs(self):
        with patch('tests._testwith.something') as mock_something:
            self.assertEqual(something, mock_something, "unpatched")
            self.assertTrue(is_instance(mock_something, MagicMock),
                            "patching wrong type")
        self.assertEqual(something, sentinel.Something)

    def testPatchObjectWithStatementAs(self):
        mock = Mock()
        original = mock.something
        with patch.object(mock, 'something'):
            self.assertNotEqual(mock.something, original, "unpatched")
        self.assertEqual(mock.something, original)

    def testWithStatementNested(self):
        with catch_warnings(record=True):
            # nested is deprecated in Python 2.7
            with nested(patch('tests._testwith.something'),
                    patch('tests._testwith.something_else')) as (mock_something, mock_something_else):
                self.assertEqual(something, mock_something, "unpatched")
                self.assertEqual(something_else, mock_something_else,
                                 "unpatched")
        self.assertEqual(something, sentinel.Something)
        self.assertEqual(something_else, sentinel.SomethingElse)

    def testWithStatementSpecified(self):
        with patch('tests._testwith.something', sentinel.Patched) as mock_something:
            self.assertEqual(something, mock_something, "unpatched")
            self.assertEqual(mock_something, sentinel.Patched, "wrong patch")
        self.assertEqual(something, sentinel.Something)

    def testContextManagerMocking(self):
        mock = Mock()
        mock.__enter__ = Mock()
        mock.__exit__ = Mock()
        mock.__exit__.return_value = False

        with mock as m:
            self.assertEqual(m, mock.__enter__.return_value)
        mock.__enter__.assert_called_with()
        mock.__exit__.assert_called_with(None, None, None)

    def testContextManagerWithMagicMock(self):
        mock = MagicMock()

        with self.assertRaises(TypeError):
            with mock:
                'foo' + 3
        mock.__enter__.assert_called_with()
        self.assertTrue(mock.__exit__.called)

    def testWithStatementSameAttribute(self):
        with patch('tests._testwith.something', sentinel.Patched) as mock_something:
            self.assertEqual(something, mock_something, "unpatched")

            with patch('tests._testwith.something') as mock_again:
                self.assertEqual(something, mock_again, "unpatched")

            self.assertEqual(something, mock_something,
                             "restored with wrong instance")

        self.assertEqual(something, sentinel.Something, "not restored")

    def testWithStatementImbricated(self):
        with patch('tests._testwith.something') as mock_something:
            self.assertEqual(something, mock_something, "unpatched")

            with patch('tests._testwith.something_else') as mock_something_else:
                self.assertEqual(something_else, mock_something_else,
                                 "unpatched")

        self.assertEqual(something, sentinel.Something)
        self.assertEqual(something_else, sentinel.SomethingElse)

    def testDictContextManager(self):
        foo = {}
        with patch.dict(foo, {'a': 'b'}):
            self.assertEqual(foo, {'a': 'b'})
        self.assertEqual(foo, {})

        with self.assertRaises(NameError):
            with patch.dict(foo, {'a': 'b'}):
                self.assertEqual(foo, {'a': 'b'})
                raise NameError('Konrad')

        self.assertEqual(foo, {})

if __name__ == '__main__':
    unittest2.main()
<commit_msg>Test name refactor<commit_after># Copyright (C) 2007-2011 Michael Foord & the mock team
# E-mail: fuzzyman AT voidspace DOT org DOT uk
# http://www.voidspace.org.uk/python/mock/

from __future__ import with_statement

from tests.support import unittest2, is_instance

from mock import MagicMock, Mock, patch, sentinel

from tests.support_with import catch_warnings, nested

something  = sentinel.Something
something_else  = sentinel.SomethingElse

class WithTest(unittest2.TestCase):

    def test_with_statement(self):
        with patch('tests._testwith.something', sentinel.Something2):
            self.assertEqual(something, sentinel.Something2, "unpatched")
        self.assertEqual(something, sentinel.Something)

    def test_with_statement_exception(self):
        try:
            with patch('tests._testwith.something', sentinel.Something2):
                self.assertEqual(something, sentinel.Something2, "unpatched")
                raise Exception('pow')
        except Exception:
            pass
        else:
            self.fail("patch swallowed exception")
        self.assertEqual(something, sentinel.Something)

    def test_with_statement_as(self):
        with patch('tests._testwith.something') as mock_something:
            self.assertEqual(something, mock_something, "unpatched")
            self.assertTrue(is_instance(mock_something, MagicMock),
                            "patching wrong type")
        self.assertEqual(something, sentinel.Something)

    def test_patch_object_with_statement(self):
        mock = Mock()
        original = mock.something
        with patch.object(mock, 'something'):
            self.assertNotEqual(mock.something, original, "unpatched")
        self.assertEqual(mock.something, original)

    def test_with_statement_nested(self):
        with catch_warnings(record=True):
            # nested is deprecated in Python 2.7
            with nested(patch('tests._testwith.something'),
                    patch('tests._testwith.something_else')) as (mock_something, mock_something_else):
                self.assertEqual(something, mock_something, "unpatched")
                self.assertEqual(something_else, mock_something_else,
                                 "unpatched")
        self.assertEqual(something, sentinel.Something)
        self.assertEqual(something_else, sentinel.SomethingElse)

    def test_with_statement_specified(self):
        with patch('tests._testwith.something', sentinel.Patched) as mock_something:
            self.assertEqual(something, mock_something, "unpatched")
            self.assertEqual(mock_something, sentinel.Patched, "wrong patch")
        self.assertEqual(something, sentinel.Something)

    def testContextManagerMocking(self):
        mock = Mock()
        mock.__enter__ = Mock()
        mock.__exit__ = Mock()
        mock.__exit__.return_value = False

        with mock as m:
            self.assertEqual(m, mock.__enter__.return_value)
        mock.__enter__.assert_called_with()
        mock.__exit__.assert_called_with(None, None, None)

    def test_context_manager_with_magic_mock(self):
        mock = MagicMock()

        with self.assertRaises(TypeError):
            with mock:
                'foo' + 3
        mock.__enter__.assert_called_with()
        self.assertTrue(mock.__exit__.called)

    def test_with_statement_same_attribute(self):
        with patch('tests._testwith.something', sentinel.Patched) as mock_something:
            self.assertEqual(something, mock_something, "unpatched")

            with patch('tests._testwith.something') as mock_again:
                self.assertEqual(something, mock_again, "unpatched")

            self.assertEqual(something, mock_something,
                             "restored with wrong instance")

        self.assertEqual(something, sentinel.Something, "not restored")

    def test_with_statement_imbricated(self):
        with patch('tests._testwith.something') as mock_something:
            self.assertEqual(something, mock_something, "unpatched")

            with patch('tests._testwith.something_else') as mock_something_else:
                self.assertEqual(something_else, mock_something_else,
                                 "unpatched")

        self.assertEqual(something, sentinel.Something)
        self.assertEqual(something_else, sentinel.SomethingElse)

    def test_dict_context_manager(self):
        foo = {}
        with patch.dict(foo, {'a': 'b'}):
            self.assertEqual(foo, {'a': 'b'})
        self.assertEqual(foo, {})

        with self.assertRaises(NameError):
            with patch.dict(foo, {'a': 'b'}):
                self.assertEqual(foo, {'a': 'b'})
                raise NameError('Konrad')

        self.assertEqual(foo, {})

if __name__ == '__main__':
    unittest2.main()
<|endoftext|>


<commit_before>#!/usr/bin/env python -u
# -*- coding: utf-8 -*-

"""
Usage:

time python -u ${PATH}/insertMECOEnergyData.py > ${LOG_FILE}

From the *current directory*, recursively descend into every existing folder and
insert all data that is found.

This script makes use of insertSingleMECOEnergyDataFile.py.

This script only supports processing of *.xml.gz files.
"""

__author__ = 'Daniel Zhang (張道博)'
__copyright__ = 'Copyright (c) 2013, University of Hawaii Smart Energy Project'
__license__ = 'https://raw.github' \
              '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD' \
              '-LICENSE.txt'

import os
import fnmatch
import sys
from subprocess import call
from msg_configer import MSGConfiger
import re
from msg_notifier import MSGNotifier
import argparse
from meco_plotting import MECOPlotting
from insertSingleMECOEnergyDataFile import Inserter
import time
from msg_logger import MSGLogger
import multiprocessing

xmlGzCount = 0
xmlCount = 0
configer = MSGConfiger()
logger = MSGLogger(__name__, 'info')
binPath = MSGConfiger.configOptionValue(configer, "Executable Paths",
                                        "bin_path")
commandLineArgs = None
global msgBody
msgBody = ''
notifier = MSGNotifier()

USE_SCRIPT_METHOD = False

def processCommandLineArguments():
    global parser, commandLineArgs
    parser = argparse.ArgumentParser(
        description = 'Perform recursive insertion of data contained in the '
                      'current directory to the MECO database specified in the '
                      'configuration file.')
    parser.add_argument('--email', action = 'store_true', default = False,
                        help = 'Send email notification if this flag is '
                               'specified.')
    parser.add_argument('--testing', action = 'store_true', default = False,
                        help = 'If this flag is on, '
                               'insert data to the testing database as '
                               'specified in the local configuration file.')
    commandLineArgs = parser.parse_args()

def makePlotAttachments():
    """
    Make data plots.
    """

    plotPath = configer.configOptionValue("Data Paths", "plot_path")
    sys.stderr.write("plotPath = %s\n" % plotPath)

    # If the plot doesn't exist then return.
    if not os.path.isdir(plotPath):
        return []

    attachments = ["%s/ReadingAndMeterCounts.png" % plotPath]
    for a in attachments:
        sys.stderr.write("attachment = %s\n" % a)
    return attachments

def logLegend():
    """
    Output a legend describing the concise report format.
    """

    legend = "Log Legend: #: = process ID, {} = dupes, () = element group, " \
             "[] = process for insert elements, <> = <reading insert count, " \
             "register insert count, event insert count, group insert count," \
             "total insert count>, * = commit\nrd = reading, re = register, " \
             "ev = event"
    return legend

def insertDataWrapper(fullPath):
    """
    A wrapper for data insertion multiprocessing.
    :returns: Log of parsing along with performance results.
    """

    pattern = 'Process-(\d+),'
    jobString = str(multiprocessing.current_process())
    match = re.search(pattern, jobString)
    assert match.group(1) is not None, "Process ID was matched."

    myLog = ''
    myLog += "\n"
    myLog += fullPath
    # print msg
    myLog += "\n"
    startTime = time.time()
    myLog += inserter.insertData(fullPath, testing = commandLineArgs.testing,
                                 jobID = match.group(1))
    myLog += "\n"
    myLog += "\nWall time = {:.2f} seconds.\n".format(
        time.time() - startTime)

    logger.log('myLog = %s' % myLog)

    return myLog

def worker(path, returnDict):
    result = insertDataWrapper(path)
    sys.stderr.write("result\n")
    sys.stderr.write(result)
    sys.stderr.write("\n")
    pattern = 'Process-(\d+),'
    jobString = str(multiprocessing.current_process())
    match = re.search(pattern, jobString)
    assert match.group(1) is not None, "Process ID was matched."
    returnDict[match.group(1)] = result

processCommandLineArguments()

inserter = Inserter()

if commandLineArgs.testing:
    logger.log("Testing mode is ON.\n", 'info')
if commandLineArgs.email:
    logger.log("Email will be sent.\n", 'info')

msg = ''
databaseName = ''

if commandLineArgs.testing:
    databaseName = configer.configOptionValue("Database", "testing_db_name")
else:
    databaseName = configer.configOptionValue("Database", "db_name")

msg = "Recursively inserting data to the database named %s." % databaseName
print msg
msgBody += msg + "\n"

startingDirectory = os.getcwd()
msg = "Starting in %s." % startingDirectory
print msg
msgBody += msg + "\n"

for root, dirnames, filenames in os.walk('.'):

    for filename in fnmatch.filter(filenames, '*.xml'):
        fullPath = os.path.join(root, filename)
        msg = fullPath
        print msg
        msgBody += msg + "\n"
        xmlCount += 1

if xmlCount != 0:
    msg = "Found XML files that are not gzip compressed.\nUnable to proceed."
    print msg
    msgBody += msg + "\n"
    if (commandLineArgs.email):
        notifier.sendNotificationEmail(msgBody, commandLineArgs.testing)
    sys.exit(-1)

insertScript = "%s/insertSingleMECOEnergyDataFile.py" % binPath
msg = "insertScript = %s" % insertScript
print msg
msgBody += msg + "\n"

parseLog = ''

try:
    with open(insertScript):
        pass
except IOError:
    msg = "Insert script %s not found." % insertScript
    print msg
    msgBody += msg + "\n"

startTime = 0

pathsToProcess = []
for root, dirnames, filenames in os.walk('.'):
    for filename in fnmatch.filter(filenames, '*.xml.gz'):
        if re.search('.*log\.xml', filename) is None: # Skip *log.xml files.
            xmlGzCount += 1
            pathsToProcess.append(os.path.join(root, filename))

try:
    pool = multiprocessing.Pool(
        int(configer.configOptionValue('Hardware', 'multiprocessing_limit')))

    procs = []
    manager = multiprocessing.Manager()
    returnDict = manager.dict()

    for path in pathsToProcess:
        procs.append(
            multiprocessing.Process(target = worker, args = (path, returnDict)))
        procs[-1].daemon = True
        procs[-1].start()

    for proc in procs:
        proc.join()

    for key in returnDict.keys():
        sys.stderr.write("Process %s results:\n" % key)
        sys.stderr.write(returnDict[key])
        sys.stderr.write("\n")
        msgBody += returnDict[key]

except Exception, e:
    logger.log('An exception occurred: %s' % e, 'error')

msgBody += "\n" + logLegend() + "\n"

msg = "\nProcessed file count is %s.\n" % xmlGzCount

print msg
msgBody += msg + "\n"

plotter = MECOPlotting(commandLineArgs.testing)

try:
    plotter.plotReadingAndMeterCounts(databaseName)
    msg = "\nPlot is attached.\n"
except:
    # @todo What exception is thrown?
    msg = "\nFailed to generate plot.\n"

msgBody += msg

if commandLineArgs.email:
    notifier.sendMailWithAttachments(msgBody, makePlotAttachments(),
                                     commandLineArgs.testing)

logger.log("msgBody = %s" % msgBody)
<commit_msg>removing unused items<commit_after>#!/usr/bin/env python -u
# -*- coding: utf-8 -*-

"""
Usage:

time python -u ${PATH}/insertMECOEnergyData.py > ${LOG_FILE}

From the *current directory*, recursively descend into every existing folder and
insert all data that is found.

This script makes use of insertSingleMECOEnergyDataFile.py.

This script only supports processing of *.xml.gz files.
"""

__author__ = 'Daniel Zhang (張道博)'
__copyright__ = 'Copyright (c) 2013, University of Hawaii Smart Energy Project'
__license__ = 'https://raw.github' \
              '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD' \
              '-LICENSE.txt'

import os
import fnmatch
import sys
from subprocess import call
from msg_configer import MSGConfiger
import re
from msg_notifier import MSGNotifier
import argparse
from meco_plotting import MECOPlotting
from insertSingleMECOEnergyDataFile import Inserter
import time
from msg_logger import MSGLogger
import multiprocessing

xmlGzCount = 0
xmlCount = 0
configer = MSGConfiger()
logger = MSGLogger(__name__, 'info')
binPath = MSGConfiger.configOptionValue(configer, "Executable Paths",
                                        "bin_path")
commandLineArgs = None
global msgBody
msgBody = ''
notifier = MSGNotifier()

USE_SCRIPT_METHOD = False

def processCommandLineArguments():
    global parser, commandLineArgs
    parser = argparse.ArgumentParser(
        description = 'Perform recursive insertion of data contained in the '
                      'current directory to the MECO database specified in the '
                      'configuration file.')
    parser.add_argument('--email', action = 'store_true', default = False,
                        help = 'Send email notification if this flag is '
                               'specified.')
    parser.add_argument('--testing', action = 'store_true', default = False,
                        help = 'If this flag is on, '
                               'insert data to the testing database as '
                               'specified in the local configuration file.')
    commandLineArgs = parser.parse_args()

def makePlotAttachments():
    """
    Make data plots.
    """

    plotPath = configer.configOptionValue("Data Paths", "plot_path")
    sys.stderr.write("plotPath = %s\n" % plotPath)

    # If the plot doesn't exist then return.
    if not os.path.isdir(plotPath):
        return []

    attachments = ["%s/ReadingAndMeterCounts.png" % plotPath]
    for a in attachments:
        sys.stderr.write("attachment = %s\n" % a)
    return attachments

def logLegend():
    """
    Output a legend describing the concise report format.
    """

    legend = "Log Legend: #: = process ID, {} = dupes, () = element group, " \
             "[] = process for insert elements, <> = <reading insert count, " \
             "register insert count, event insert count, group insert count," \
             "total insert count>, * = commit\nrd = reading, re = register, " \
             "ev = event"
    return legend

def insertDataWrapper(fullPath):
    """
    A wrapper for data insertion multiprocessing.
    :returns: Log of parsing along with performance results.
    """

    pattern = 'Process-(\d+),'
    jobString = str(multiprocessing.current_process())
    match = re.search(pattern, jobString)
    assert match.group(1) is not None, "Process ID was matched."

    myLog = ''
    myLog += "\n"
    myLog += fullPath
    # print msg
    myLog += "\n"
    startTime = time.time()
    myLog += inserter.insertData(fullPath, testing = commandLineArgs.testing,
                                 jobID = match.group(1))
    myLog += "\n"
    myLog += "\nWall time = {:.2f} seconds.\n".format(
        time.time() - startTime)

    logger.log('myLog = %s' % myLog)

    return myLog

def worker(path, returnDict):
    result = insertDataWrapper(path)
    # sys.stderr.write("result\n")
    # sys.stderr.write(result)
    # sys.stderr.write("\n")
    pattern = 'Process-(\d+),'
    jobString = str(multiprocessing.current_process())
    match = re.search(pattern, jobString)
    assert match.group(1) is not None, "Process ID was matched."
    returnDict[match.group(1)] = result

processCommandLineArguments()

inserter = Inserter()

if commandLineArgs.testing:
    logger.log("Testing mode is ON.\n", 'info')
if commandLineArgs.email:
    logger.log("Email will be sent.\n", 'info')

msg = ''
databaseName = ''

if commandLineArgs.testing:
    databaseName = configer.configOptionValue("Database", "testing_db_name")
else:
    databaseName = configer.configOptionValue("Database", "db_name")

msg = "Recursively inserting data to the database named %s." % databaseName
print msg
msgBody += msg + "\n"

startingDirectory = os.getcwd()
msg = "Starting in %s." % startingDirectory
print msg
msgBody += msg + "\n"

for root, dirnames, filenames in os.walk('.'):

    for filename in fnmatch.filter(filenames, '*.xml'):
        fullPath = os.path.join(root, filename)
        msg = fullPath
        print msg
        msgBody += msg + "\n"
        xmlCount += 1

if xmlCount != 0:
    msg = "Found XML files that are not gzip compressed.\nUnable to proceed."
    print msg
    msgBody += msg + "\n"
    if (commandLineArgs.email):
        notifier.sendNotificationEmail(msgBody, commandLineArgs.testing)
    sys.exit(-1)

insertScript = "%s/insertSingleMECOEnergyDataFile.py" % binPath
msg = "insertScript = %s" % insertScript
print msg
msgBody += msg + "\n"

parseLog = ''

try:
    with open(insertScript):
        pass
except IOError:
    msg = "Insert script %s not found." % insertScript
    print msg
    msgBody += msg + "\n"

startTime = 0

pathsToProcess = []
for root, dirnames, filenames in os.walk('.'):
    for filename in fnmatch.filter(filenames, '*.xml.gz'):
        if re.search('.*log\.xml', filename) is None: # Skip *log.xml files.
            xmlGzCount += 1
            pathsToProcess.append(os.path.join(root, filename))

try:
    # pool = multiprocessing.Pool(
    #     int(configer.configOptionValue('Hardware', 'multiprocessing_limit')))

    procs = []
    manager = multiprocessing.Manager()
    returnDict = manager.dict()

    for path in pathsToProcess:
        procs.append(
            multiprocessing.Process(target = worker, args = (path, returnDict)))
        procs[-1].daemon = True
        procs[-1].start()

    for proc in procs:
        proc.join()

    for key in returnDict.keys():
        sys.stderr.write("Process %s results:\n" % key)
        sys.stderr.write(returnDict[key])
        sys.stderr.write("\n")
        msgBody += returnDict[key]

except Exception, e:
    logger.log('An exception occurred: %s' % e, 'error')

msgBody += "\n" + logLegend() + "\n"

msg = "\nProcessed file count is %s.\n" % xmlGzCount

print msg
msgBody += msg + "\n"

plotter = MECOPlotting(commandLineArgs.testing)

try:
    plotter.plotReadingAndMeterCounts(databaseName)
    msg = "\nPlot is attached.\n"
except Exception, e:
    msg = "\nFailed to generate plot.\n"
    logger.log('An exception occurred: Failed to generate plot. %s' % e,
               'error')

msgBody += msg

if commandLineArgs.email:
    notifier.sendMailWithAttachments(msgBody, makePlotAttachments(),
                                     commandLineArgs.testing)

logger.log("msgBody = %s" % msgBody)
<|endoftext|>


<commit_before>##########################################################################
#
# Copyright 2008-2009 VMware, Inc.
# All Rights Reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
##########################################################################/

"""GLX API description."""

from stdapi import *
from glapi import glapi

VisualID = Alias("VisualID", UInt32)
Display = Opaque("Display *")
Visual = Opaque("Visual *")
Font = Alias("Font", UInt32)
Pixmap = Alias("Pixmap", UInt32)
Window = Alias("Window", UInt32)

GLXContext = Opaque("GLXContext")
GLXPixmap = Alias("GLXPixmap", UInt32)
GLXDrawable = Alias("GLXDrawable", UInt32)
GLXFBConfig = Opaque("GLXFBConfig")
GLXFBConfigID = Alias("GLXFBConfigID", UInt32)
GLXContextID = Alias("GLXContextID", UInt32)
GLXWindow = Alias("GLXWindow", UInt32)
GLXPbuffer = Alias("GLXPbuffer", UInt32)

XVisualInfo = Struct("XVisualInfo", [
  (Visual, "visual"),
  (VisualID, "visualid"),
  (Int, "screen"),
  (Int, "depth"),
  (Int, "c_class"),
  (ULong, "red_mask"),
  (ULong, "green_mask"),
  (ULong, "blue_mask"),
  (Int, "colormap_size"),
  (Int, "bits_per_rgb"),
])

Bool_ = Alias("Bool", Int)

glxapi = API("GLX")

PROC = Opaque("__GLXextFuncPtr")

glxapi.add_functions([
    # GLX
    Function(Pointer(XVisualInfo), "glXChooseVisual", [(Display, "dpy"), (Int, "screen"), (Pointer(Int), "attribList")]),
    Function(GLXContext, "glXCreateContext", [(Display, "dpy"), (Pointer(XVisualInfo), "vis"), (GLXContext, "shareList"), (Bool_, "direct")]),
    Function(Void, "glXDestroyContext",  [(Display, "dpy"), (GLXContext, "ctx")]),
    Function(Bool_, "glXMakeCurrent", [(Display, "dpy"), (GLXDrawable, "drawable"), (GLXContext, "ctx")]),
    Function(Void, "glXCopyContext", [(Display, "dpy"), (GLXContext, "src"), (GLXContext, "dst"),
                                      (ULong, "mask")]),
    Function(Void, "glXSwapBuffers", [(Display, "dpy"), (GLXDrawable, "drawable")]),
    Function(GLXPixmap, "glXCreateGLXPixmap", [(Display, "dpy"), (Pointer(XVisualInfo), "visual"),
                                               (Pixmap, "pixmap")]),
    Function(Void, "glXDestroyGLXPixmap", [(Display, "dpy"), (GLXPixmap, "pixmap")]),
    Function(Bool_, "glXQueryExtension", [(Display, "dpy"), (Pointer(Int), "errorb"), (Pointer(Int), "event")]),
    Function(Bool_, "glXQueryVersion", [(Display, "dpy"), (Pointer(Int), "maj"), (Pointer(Int), "min")]),
    Function(Bool_, "glXIsDirect", [(Display, "dpy"), (GLXContext, "ctx")]),
    Function(Int, "glXGetConfig", [(Display, "dpy"), (Pointer(XVisualInfo), "visual"),
                                    (Int, "attrib"), (Pointer(Int), "value")]),
    Function(GLXContext, "glXGetCurrentContext", [], sideeffects=False),
    Function(GLXDrawable, "glXGetCurrentDrawable", [], sideeffects=False),
    Function(Void, "glXWaitGL", []),
    Function(Void, "glXWaitX", []),
    Function(Void, "glXUseXFont", [(Font, "font"), (Int, "first"), (Int, "count"), (Int, "list")]),

    # GLX 1.1 and later
    Function((Const(String("char *"))), "glXQueryExtensionsString", [(Display, "dpy"), (Int, "screen")]),
    Function((Const(String("char *"))), "glXQueryServerString",  [(Display, "dpy"), (Int, "screen"), (Int, "name")]),
    Function((Const(String("char *"))), "glXGetClientString", [(Display, "dpy"), (Int, "name")]),

    # GLX 1.2 and later
    Function(Display, "glXGetCurrentDisplay", [], sideeffects=False),

    # GLX 1.3 and later
    Function(Pointer(GLXFBConfig), "glXChooseFBConfig", [(Display, "dpy"), (Int, "screen"),
                                                         (Pointer(Const(Int)), "attribList"), (Pointer(Int), "nitems")]),
    Function(Int, "glXGetFBConfigAttrib", [(Display, "dpy"), (GLXFBConfig, "config"),
                                            (Int, "attribute"), (Pointer(Int), "value")]),
    Function(Pointer(GLXFBConfig), "glXGetFBConfigs", [(Display, "dpy"), (Int, "screen"),
                                              (Pointer(Int), "nelements")]),
    Function(Pointer(XVisualInfo), "glXGetVisualFromFBConfig", [(Display, "dpy"),
                                                                (GLXFBConfig, "config")]),
    Function(GLXWindow, "glXCreateWindow", [(Display, "dpy"), (GLXFBConfig, "config"),
                                            (Window, "win"), (Pointer(Const(Int)), "attribList")]),
    Function(Void, "glXDestroyWindow", [(Display, "dpy"), (GLXWindow, "window")]),
    Function(GLXPixmap, "glXCreatePixmap", [(Display, "dpy"), (GLXFBConfig, "config"),
                                            (Pixmap, "pixmap"), (Pointer(Const(Int)), "attribList")]),
    Function(Void, "glXDestroyPixmap", [(Display, "dpy"), (GLXPixmap, "pixmap")]),
    Function(GLXPbuffer, "glXCreatePbuffer", [(Display, "dpy"), (GLXFBConfig, "config"),
                                              (Pointer(Const(Int)), "attribList")]),
    Function(Void, "glXDestroyPbuffer", [(Display, "dpy"), (GLXPbuffer, "pbuf")]),
    Function(Void, "glXQueryDrawable", [(Display, "dpy"), (GLXDrawable, "draw"), (Int, "attribute"),
                                        (Pointer(UInt), "value")]),
    Function(GLXContext, "glXCreateNewContext", [(Display, "dpy"), (GLXFBConfig, "config"),
                                                 (Int, "renderType"), (GLXContext, "shareList"),
                                                 (Bool_, "direct")]),
    Function(Bool_, "glXMakeContextCurrent", [(Display, "dpy"), (GLXDrawable, "draw"),
                                             (GLXDrawable, "read"), (GLXContext, "ctx")]),
    Function(GLXDrawable, "glXGetCurrentReadDrawable", []),
    Function(Int, "glXQueryContext", [(Display, "dpy"), (GLXContext, "ctx"), (Int, "attribute"),
                                      (Pointer(Int), "value")]),
    Function(Void, "glXSelectEvent", [(Display, "dpy"), (GLXDrawable, "drawable"),
                                      (ULong, "mask")]),
    Function(Void, "glXGetSelectedEvent", [(Display, "dpy"), (GLXDrawable, "drawable"),
                                           (Pointer(ULong), "mask")]),
    # TODO: extensions

    # Must be last
    Function(PROC, "glXGetProcAddressARB", [(Alias("const GLubyte *", CString), "procName")]),
    Function(PROC, "glXGetProcAddress", [(Alias("const GLubyte *", CString), "procName")]),
])

<commit_msg>More detail in glxapi.<commit_after>##########################################################################
#
# Copyright 2008-2009 VMware, Inc.
# All Rights Reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
##########################################################################/

"""GLX API description."""

from stdapi import *
from glapi import glapi

VisualID = Alias("VisualID", UInt32)
Display = Opaque("Display *")
Visual = Opaque("Visual *")
Font = Alias("Font", UInt32)
Pixmap = Alias("Pixmap", UInt32)
Window = Alias("Window", UInt32)

GLXContext = Opaque("GLXContext")
GLXPixmap = Alias("GLXPixmap", UInt32)
GLXDrawable = Alias("GLXDrawable", UInt32)
GLXFBConfig = Opaque("GLXFBConfig")
GLXFBConfigID = Alias("GLXFBConfigID", UInt32)
GLXContextID = Alias("GLXContextID", UInt32)
GLXWindow = Alias("GLXWindow", UInt32)
GLXPbuffer = Alias("GLXPbuffer", UInt32)

XVisualInfo = Struct("XVisualInfo", [
  (Visual, "visual"),
  (VisualID, "visualid"),
  (Int, "screen"),
  (Int, "depth"),
  (Int, "c_class"),
  (ULong, "red_mask"),
  (ULong, "green_mask"),
  (ULong, "blue_mask"),
  (Int, "colormap_size"),
  (Int, "bits_per_rgb"),
])

Bool_ = Enum("Bool", [
    "False",
    "True",
])

GLXAttrib = FakeEnum(Int, [
    "GLX_USE_GL",
    "GLX_BUFFER_SIZE",
    "GLX_LEVEL",
    "GLX_RGBA",
    "GLX_DOUBLEBUFFER",
    "GLX_STEREO",
    "GLX_AUX_BUFFERS",
    "GLX_RED_SIZE",
    "GLX_GREEN_SIZE",
    "GLX_BLUE_SIZE",
    "GLX_ALPHA_SIZE",
    "GLX_DEPTH_SIZE",
    "GLX_STENCIL_SIZE",
    "GLX_ACCUM_RED_SIZE",
    "GLX_ACCUM_GREEN_SIZE",
    "GLX_ACCUM_BLUE_SIZE",
    "GLX_ACCUM_ALPHA_SIZE",
    "GLX_CONFIG_CAVEAT",
    "GLX_X_VISUAL_TYPE",
    "GLX_TRANSPARENT_TYPE",
    "GLX_TRANSPARENT_INDEX_VALUE",
    "GLX_TRANSPARENT_RED_VALUE",
    "GLX_TRANSPARENT_GREEN_VALUE",
    "GLX_TRANSPARENT_BLUE_VALUE",
    "GLX_TRANSPARENT_ALPHA_VALUE",
    "GLX_DONT_CARE",
])

GLXError = FakeEnum(Int, [
    "GLX_BAD_SCREEN",
    "GLX_BAD_ATTRIBUTE",
    "GLX_NO_EXTENSION",
    "GLX_BAD_VISUAL",
    "GLX_BAD_CONTEXT",
    "GLX_BAD_VALUE",
    "GLX_BAD_ENUM",
])

GLXbuffer = Flags(Int, [
    "GLX_WINDOW_BIT",
    "GLX_PIXMAP_BIT",
    "GLX_PBUFFER_BIT",
    "GLX_AUX_BUFFERS_BIT",
    "GLX_FRONT_LEFT_BUFFER_BIT",
    "GLX_FRONT_RIGHT_BUFFER_BIT",
    "GLX_BACK_LEFT_BUFFER_BIT",
    "GLX_BACK_RIGHT_BUFFER_BIT",
    "GLX_DEPTH_BUFFER_BIT",
    "GLX_STENCIL_BUFFER_BIT",
    "GLX_ACCUM_BUFFER_BIT",
])

GLXEnum = FakeEnum(Int, [
    "GLX_NONE",
    "GLX_SLOW_CONFIG",
    "GLX_TRUE_COLOR",
    "GLX_DIRECT_COLOR",
    "GLX_PSEUDO_COLOR",
    "GLX_STATIC_COLOR",
    "GLX_GRAY_SCALE",
    "GLX_STATIC_GRAY",
    "GLX_TRANSPARENT_RGB",
    "GLX_TRANSPARENT_INDEX",
    "GLX_VISUAL_ID",
    "GLX_SCREEN",
    "GLX_NON_CONFORMANT_CONFIG",
    "GLX_DRAWABLE_TYPE",
    "GLX_RENDER_TYPE",
    "GLX_X_RENDERABLE",
    "GLX_FBCONFIG_ID",
    "GLX_RGBA_TYPE",
    "GLX_COLOR_INDEX_TYPE",
    "GLX_MAX_PBUFFER_WIDTH",
    "GLX_MAX_PBUFFER_HEIGHT",
    "GLX_MAX_PBUFFER_PIXELS",
    "GLX_PRESERVED_CONTENTS",
    "GLX_LARGEST_PBUFFER",
    "GLX_WIDTH",
    "GLX_HEIGHT",
    "GLX_EVENT_MASK",
    "GLX_DAMAGED",
    "GLX_SAVED",
    "GLX_WINDOW",
    "GLX_PBUFFER",
    "GLX_PBUFFER_HEIGHT",
    "GLX_PBUFFER_WIDTH",
])

GLXbuffer = Flags(Int, [
    "GLX_RGBA_BIT",
    "GLX_COLOR_INDEX_BIT",
    "GLX_PBUFFER_CLOBBER_MASK",
])

glxapi = API("GLX")

PROC = Opaque("__GLXextFuncPtr")

glxapi.add_functions([
    # GLX
    Function(Pointer(XVisualInfo), "glXChooseVisual", [(Display, "dpy"), (Int, "screen"), (Array(GLXAttrib, "__AttribList_size(attribList)"), "attribList")]),
    Function(GLXContext, "glXCreateContext", [(Display, "dpy"), (Pointer(XVisualInfo), "vis"), (GLXContext, "shareList"), (Bool_, "direct")]),
    Function(Void, "glXDestroyContext",  [(Display, "dpy"), (GLXContext, "ctx")]),
    Function(Bool_, "glXMakeCurrent", [(Display, "dpy"), (GLXDrawable, "drawable"), (GLXContext, "ctx")]),
    Function(Void, "glXCopyContext", [(Display, "dpy"), (GLXContext, "src"), (GLXContext, "dst"),
                                      (ULong, "mask")]),
    Function(Void, "glXSwapBuffers", [(Display, "dpy"), (GLXDrawable, "drawable")]),
    Function(GLXPixmap, "glXCreateGLXPixmap", [(Display, "dpy"), (Pointer(XVisualInfo), "visual"),
                                               (Pixmap, "pixmap")]),
    Function(Void, "glXDestroyGLXPixmap", [(Display, "dpy"), (GLXPixmap, "pixmap")]),
    Function(Bool_, "glXQueryExtension", [(Display, "dpy"), Out(Pointer(Int), "errorb"), Out(Pointer(Int), "event")]),
    Function(Bool_, "glXQueryVersion", [(Display, "dpy"), Out(Pointer(Int), "maj"), Out(Pointer(Int), "min")]),
    Function(Bool_, "glXIsDirect", [(Display, "dpy"), (GLXContext, "ctx")]),
    Function(GLXError, "glXGetConfig", [(Display, "dpy"), (Pointer(XVisualInfo), "visual"),
                                    (GLXAttrib, "attrib"), Out(Pointer(Int), "value")]),
    Function(GLXContext, "glXGetCurrentContext", [], sideeffects=False),
    Function(GLXDrawable, "glXGetCurrentDrawable", [], sideeffects=False),
    Function(Void, "glXWaitGL", []),
    Function(Void, "glXWaitX", []),
    Function(Void, "glXUseXFont", [(Font, "font"), (Int, "first"), (Int, "count"), (Int, "list")]),

    # GLX 1.1 and later
    Function((Const(String("char *"))), "glXQueryExtensionsString", [(Display, "dpy"), (Int, "screen")]),
    Function((Const(String("char *"))), "glXQueryServerString",  [(Display, "dpy"), (Int, "screen"), (Int, "name")]),
    Function((Const(String("char *"))), "glXGetClientString", [(Display, "dpy"), (Int, "name")]),

    # GLX 1.2 and later
    Function(Display, "glXGetCurrentDisplay", [], sideeffects=False),

    # GLX 1.3 and later
    Function(Array(GLXFBConfig, "*nitems"), "glXChooseFBConfig", [(Display, "dpy"), (Int, "screen"), (Const(Array(GLXAttrib, "__AttribList_size(attribList)")), "attribList"), Out(Pointer(Int), "nitems")]),
    Function(Int, "glXGetFBConfigAttrib", [(Display, "dpy"), (GLXFBConfig, "config"), (GLXAttrib, "attribute"), Out(Pointer(Int), "value")]),
    Function(Array(GLXFBConfig, "*nelements"), "glXGetFBConfigs", [(Display, "dpy"), (Int, "screen"),
                                                                   Out(Pointer(Int), "nelements")]),
    Function(Pointer(XVisualInfo), "glXGetVisualFromFBConfig", [(Display, "dpy"),
                                                                (GLXFBConfig, "config")]),
    Function(GLXWindow, "glXCreateWindow", [(Display, "dpy"), (GLXFBConfig, "config"),
                                            (Window, "win"), (Const(Array(Int, "__AttribList_size(attribList)")), "attribList")]),
    Function(Void, "glXDestroyWindow", [(Display, "dpy"), (GLXWindow, "window")]),
    Function(GLXPixmap, "glXCreatePixmap", [(Display, "dpy"), (GLXFBConfig, "config"),
                                            (Pixmap, "pixmap"), (Const(Array(Int, "__AttribList_size(attribList)")), "attribList")]),
    Function(Void, "glXDestroyPixmap", [(Display, "dpy"), (GLXPixmap, "pixmap")]),
    Function(GLXPbuffer, "glXCreatePbuffer", [(Display, "dpy"), (GLXFBConfig, "config"),
                                              (Const(Array(GLXEnum, "__AttribList_size(attribList)")), "attribList")]),
    Function(Void, "glXDestroyPbuffer", [(Display, "dpy"), (GLXPbuffer, "pbuf")]),
    Function(Void, "glXQueryDrawable", [(Display, "dpy"), (GLXDrawable, "draw"), (GLXEnum, "attribute"),
                                        Out(Pointer(UInt), "value")]),
    Function(GLXContext, "glXCreateNewContext", [(Display, "dpy"), (GLXFBConfig, "config"),
                                                 (GLXEnum, "renderType"), (GLXContext, "shareList"),
                                                 (Bool_, "direct")]),
    Function(Bool_, "glXMakeContextCurrent", [(Display, "dpy"), (GLXDrawable, "draw"),
                                             (GLXDrawable, "read"), (GLXContext, "ctx")]),
    Function(GLXDrawable, "glXGetCurrentReadDrawable", []),
    Function(Int, "glXQueryContext", [(Display, "dpy"), (GLXContext, "ctx"), (GLXEnum, "attribute"),
                                      Out(Pointer(Int), "value")]),
    Function(Void, "glXSelectEvent", [(Display, "dpy"), (GLXDrawable, "drawable"),
                                      (ULong, "mask")]),
    Function(Void, "glXGetSelectedEvent", [(Display, "dpy"), (GLXDrawable, "drawable"),
                                           Out(Pointer(ULong), "mask")]),
    # TODO: extensions

    # Must be last
    Function(PROC, "glXGetProcAddressARB", [(Alias("const GLubyte *", CString), "procName")]),
    Function(PROC, "glXGetProcAddress", [(Alias("const GLubyte *", CString), "procName")]),
])

<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""
Model

Automatization of MitM Attack on WiFi Networks
Bachelor's Thesis UIFS FIT VUT
Martin Vondracek
2016
"""
import logging
import netifaces
import os
import re
import shutil
import subprocess

__author__ = 'Martin Vondracek'
__email__ = 'xvondr20@stud.fit.vutbr.cz'

logger = logging.getLogger(__name__)

class WirelessStation(object):
    def __str__(self, *args, **kwargs):  # TODO (xvondr20) just for debugging
        return 'WirelessStation(' + ', '.join([
            self.mac_address,
            self.power
        ]) + ')'

    def __init__(self, mac_address, power):
        self.mac_address = mac_address
        self.power = power

        self.associated_ap = None

class WirelessAccessPoint(object):
    def __str__(self, *args, **kwargs):  # TODO (xvondr20) just for debugging
        s = 'WirelessAccessPoint(' + ', '.join([
            self.essid,
            self.bssid
        ])

        if self.is_cracked():
            if 'WEP' in self.encryption:
                s += ', PSK(0x' + self.cracked_psk + ', "' + bytes.fromhex(self.cracked_psk).decode('ascii') + '"), '
            else:
                s += ', PSK("' + self.cracked_psk + '"), '

        s += ', '.join([
            self.power,
            self.channel,
            self.encryption,
            self.cipher,
            self.authentication,
            str(self.wps),
            self.iv_sum
        ]) + ')'
        return s

    def __init__(self, bssid, power, channel, encryption, cipher, authentication, wps, essid, iv_sum):
        self.bssid = bssid
        self.power = power
        self.channel = channel
        self.encryption = encryption
        self.cipher = cipher
        self.authentication = authentication
        self.wps = wps
        self.essid = essid
        self.iv_sum = iv_sum

        self.associated_stations = list()

        self.arp_cap_path = None
        """path to the capture of ARP Requests, if available"""
        self.psk_path = None
        """path to the file containing hexadecimal PSK, if available"""
        self.prga_xor_path = None
        """path to the file containing PRGA XOR keystream, if available"""
        self.wpa_handshake_cap_path = None
        """path to the capture of WPA handshake, if available"""

        # default paths
        self.default_arp_cap_path = os.path.join(self.dir_path, 'ARP.cap')
        if 'WEP' in self.encryption:
            self.default_psk_path = os.path.join(self.dir_path, self.encryption + '_PSK.hex')
        else:
            self.default_psk_path = os.path.join(self.dir_path, self.encryption + '_PSK.txt')
        self.default_prga_xor_path = os.path.join(self.dir_path, 'PRGA.xor')
        self.default_wpa_handshake_cap_path = os.path.join(self.dir_path, 'WPA_handshake.cap')

    @property
    def dir_path(self):
        """
        Get path to directory which should be dedicated for files related to this network.
        It the directory does not exist, the attacker is responsible for its creation using `self.make_dir()`.
        :return: str
        """
        return os.path.join(os.getcwd(), 'networks', self.essid)  # TODO (xvondr20) what is essid is not available?

    def make_dir(self):
        """
        Make directory to store files related to this network, if does not exist already.
        """
        os.makedirs(self.dir_path, exist_ok=True)

    def is_cracked(self):
        """
        Decide whether the network have been successfully cracked and therefore a PSK is available.
        :return: bool
        """
        return self.psk_path is not None  # TODO(xvondr20) WPS?

    @property
    def cracked_psk(self):
        """
        Get cracked PSK if available.
        If encryption is WEP, PSK is hexadecimal sequence. If encryption is WPA or WPA2, PSK is ASCII sequence.
        If the network have not been cracked yet, therefore PSK is not available, returns None.
        :return: str|None
        """
        if self.psk_path:
            with open(self.psk_path, 'r') as f:
                return f.read()

    def save_arp_cap(self, source_arp_cap_path):
        """
        Save capture with ARP Requests for successful ARP Replay.
        Overwrites previous file, if any exists.
        :param source_arp_cap_path: path to capture of ARP Requests
        """
        if not os.path.isfile(source_arp_cap_path):
            raise FileNotFoundError
        shutil.move(source_arp_cap_path, self.default_arp_cap_path)
        self.arp_cap_path = self.default_arp_cap_path

    def save_psk_file(self, source_psk_file_path):
        """
        Save PSK file containing hexadecimal cracked key for network.
        Overwrites previous file, if any exists.
        :param source_psk_file_path: path to PSK file
        """
        if not os.path.isfile(source_psk_file_path):
            raise FileNotFoundError
        shutil.move(source_psk_file_path, self.default_psk_path)
        self.psk_path = self.default_psk_path

    def save_wpa_handshake_cap(self, source_wpa_handshake_cap_path):
        """
        Save capture with WPA handshake.
        Overwrites previous file, if any exists.
        :param source_wpa_handshake_cap_path: path to capture with WPA handshake
        """
        if not os.path.isfile(source_wpa_handshake_cap_path):
            raise FileNotFoundError
        shutil.move(source_wpa_handshake_cap_path, self.default_wpa_handshake_cap_path)
        self.wpa_handshake_cap_path = self.default_wpa_handshake_cap_path

    def save_prga_xor(self, source_prga_xor_path):
        """
        Save file containing PRGA XOR keystream.
        Overwrites previous file, if any exists.
        :param source_prga_xor_path: path to file containing PRGA XOR keystream
        """
        if not os.path.isfile(source_prga_xor_path):
            raise FileNotFoundError
        shutil.move(source_prga_xor_path, self.default_prga_xor_path)
        self.prga_xor_path = self.default_prga_xor_path

    def add_associated_station(self, station):
        station.associated_ap = self
        self.associated_stations.append(station)

    def update_known(self):
        """
        Update known network based on saved files.
        """
        if not self.arp_cap_path and os.path.isfile(self.default_arp_cap_path):
            self.arp_cap_path = self.default_arp_cap_path
            logger.debug(self.essid + ' arp_cap known')

        if not self.psk_path and os.path.isfile(self.default_psk_path):
            self.psk_path = self.default_psk_path
            logger.debug(self.essid + ' psk known')

        if not self.prga_xor_path and os.path.isfile(self.default_prga_xor_path):
            self.prga_xor_path = self.default_prga_xor_path
            logger.debug(self.essid + ' prga_xor known')

        if not self.wpa_handshake_cap_path and os.path.isfile(self.default_wpa_handshake_cap_path):
            self.wpa_handshake_cap_path = self.default_wpa_handshake_cap_path
            logger.debug(self.essid + ' wpa_handshake_cap known')

class WirelessInterface(object):
    def __str__(self, *args, **kwargs):  # TODO (xvondr20) just for debugging
        s = 'WirelessInterface(' + ', '.join([
            self.name,
            self.mac_address,
            self.channel,
            self.driver,
            self.chipset
        ])
        if self.monitor_mode:
            s += ', monitor'
        s += ')'
        return s

    def __init__(self, name, driver=None, chipset=None):
        """
        Raises:
            ValueError if name is not a valid interface name
        """
        self.name_original = name
        self.name_monitor = None

        # get MAC address
        self.mac_address_original = self.get_mac_by_name(self.name)
        self.mac_address_spoofed = None

        self.channel = None
        self.monitor_mode = False

        # additional data
        self.driver = driver
        self.chipset = chipset

    @staticmethod
    def get_wireless_interface_obj(interface):
        """
        Get WirelessInterface object based on provided argument.
        If interface is already WirelessInterface object, it is just returned. If interface is a valid interface name,
        appropriate WirelessInterface object is created and returned.
        :param interface: WirelessInterface object or string representing valid wireless interface name
        :return: WirelessInterface object
        Raises:
            ValueError if provided interface string is not a valid interface name
            TypeError if provided interface is not a string nor a WirelessInterface object
        """
        if isinstance(interface, WirelessInterface):
            return interface
        elif isinstance(interface, str):
            return WirelessInterface(name=interface)
        else:
            raise TypeError

    @staticmethod
    def get_mac_by_name(name):
        """
        Get MAC address of interface specified by name of the interface.
        :param name: name of the network interface
        :return: string MAC address
        """
        # TODO(xvondr20) Is this safe?
        ifa = netifaces.ifaddresses(name)
        mac = ifa[netifaces.AF_LINK][0]['addr']
        return mac

    @property
    def mac_address(self):
        """
        Get current MAC address.
        """
        assert self.mac_address_spoofed or self.mac_address_original, 'No MAC address available.'

        if self.mac_address_spoofed:
            return self.mac_address_spoofed
        else:
            return self.mac_address_original

    @property
    def name(self):
        """
        Get current interface name.
        """
        assert self.name_monitor or self.name_original, 'No interface name available.'

        if self.name_monitor:
            return self.name_monitor
        else:
            return self.name_original

    @property
    def gateway(self):
        """
        Get current default gateway.
        """
        gateway = None
        gws = netifaces.gateways()
        for gw in gws[netifaces.AF_INET]:
            if gw[1] == self.name:
                gateway = gw[0]
                break
        assert gateway, 'No default gateway available.'
        return gateway

    def start_monitor_mode(self, channel=None):
        """
        :param channel: monitor interface channel
        Raises:
            CalledProcessError if airmon-ng returncode is non-zero
        """
        assert not self.monitor_mode, 'Interface already in monitor mode.'

        cre_mon_enabled = re.compile(
            r'^\s+\(\S+ monitor mode vif enabled for \[\S+\](?P<name>\S+) on \[\S+\](?P<mon>\S+)\)$')

        cmd = ['airmon-ng', 'start', self.name]
        if channel:
            cmd.append(str(channel))
            self.channel = channel
        process = subprocess.run(cmd,
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                 universal_newlines=True)

        process.check_returncode()
        # check stderr
        # TODO (xvondr20) Does 'airmon-ng' ever print anything to stderr?
        assert process.stderr == ''

        for line in process.stdout.splitlines():
            m = cre_mon_enabled.match(line)
            if m:
                self.monitor_mode = True
                self.name_monitor = m.group('mon')
                break

    def stop_monitor_mode(self):
        """
        Raises:
            CalledProcessError if airmon-ng returncode is non-zero
        """
        assert self.monitor_mode, 'Interface not in monitor mode.'

        cre_mon_disabled = re.compile(r'^\s+\(\S+ monitor mode vif disabled for \[\S+\](?P<mon>\S+)\)$')

        cmd = ['airmon-ng', 'stop', self.name]
        process = subprocess.run(cmd,
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                 universal_newlines=True)

        process.check_returncode()
        # check stderr
        # TODO (xvondr20) Does 'airmon-ng' ever print anything to stderr?
        assert process.stderr == ''

        for line in process.stdout.splitlines():
            m = cre_mon_disabled.match(line)
            if m:
                self.monitor_mode = False
                self.name_monitor = None
                break
<commit_msg>added interface_exists<commit_after>#!/usr/bin/env python3
"""
Model

Automatization of MitM Attack on WiFi Networks
Bachelor's Thesis UIFS FIT VUT
Martin Vondracek
2016
"""
import logging
import netifaces
import os
import re
import shutil
import subprocess

__author__ = 'Martin Vondracek'
__email__ = 'xvondr20@stud.fit.vutbr.cz'

logger = logging.getLogger(__name__)

class WirelessStation(object):
    def __str__(self, *args, **kwargs):  # TODO (xvondr20) just for debugging
        return 'WirelessStation(' + ', '.join([
            self.mac_address,
            self.power
        ]) + ')'

    def __init__(self, mac_address, power):
        self.mac_address = mac_address
        self.power = power

        self.associated_ap = None

class WirelessAccessPoint(object):
    def __str__(self, *args, **kwargs):  # TODO (xvondr20) just for debugging
        s = 'WirelessAccessPoint(' + ', '.join([
            self.essid,
            self.bssid
        ])

        if self.is_cracked():
            if 'WEP' in self.encryption:
                s += ', PSK(0x' + self.cracked_psk + ', "' + bytes.fromhex(self.cracked_psk).decode('ascii') + '"), '
            else:
                s += ', PSK("' + self.cracked_psk + '"), '

        s += ', '.join([
            self.power,
            self.channel,
            self.encryption,
            self.cipher,
            self.authentication,
            str(self.wps),
            self.iv_sum
        ]) + ')'
        return s

    def __init__(self, bssid, power, channel, encryption, cipher, authentication, wps, essid, iv_sum):
        self.bssid = bssid
        self.power = power
        self.channel = channel
        self.encryption = encryption
        self.cipher = cipher
        self.authentication = authentication
        self.wps = wps
        self.essid = essid
        self.iv_sum = iv_sum

        self.associated_stations = list()

        self.arp_cap_path = None
        """path to the capture of ARP Requests, if available"""
        self.psk_path = None
        """path to the file containing hexadecimal PSK, if available"""
        self.prga_xor_path = None
        """path to the file containing PRGA XOR keystream, if available"""
        self.wpa_handshake_cap_path = None
        """path to the capture of WPA handshake, if available"""

        # default paths
        self.default_arp_cap_path = os.path.join(self.dir_path, 'ARP.cap')
        if 'WEP' in self.encryption:
            self.default_psk_path = os.path.join(self.dir_path, self.encryption + '_PSK.hex')
        else:
            self.default_psk_path = os.path.join(self.dir_path, self.encryption + '_PSK.txt')
        self.default_prga_xor_path = os.path.join(self.dir_path, 'PRGA.xor')
        self.default_wpa_handshake_cap_path = os.path.join(self.dir_path, 'WPA_handshake.cap')

    @property
    def dir_path(self):
        """
        Get path to directory which should be dedicated for files related to this network.
        It the directory does not exist, the attacker is responsible for its creation using `self.make_dir()`.
        :return: str
        """
        return os.path.join(os.getcwd(), 'networks', self.essid)  # TODO (xvondr20) what is essid is not available?

    def make_dir(self):
        """
        Make directory to store files related to this network, if does not exist already.
        """
        os.makedirs(self.dir_path, exist_ok=True)

    def is_cracked(self):
        """
        Decide whether the network have been successfully cracked and therefore a PSK is available.
        :return: bool
        """
        return self.psk_path is not None  # TODO(xvondr20) WPS?

    @property
    def cracked_psk(self):
        """
        Get cracked PSK if available.
        If encryption is WEP, PSK is hexadecimal sequence. If encryption is WPA or WPA2, PSK is ASCII sequence.
        If the network have not been cracked yet, therefore PSK is not available, returns None.
        :return: str|None
        """
        if self.psk_path:
            with open(self.psk_path, 'r') as f:
                return f.read()

    def save_arp_cap(self, source_arp_cap_path):
        """
        Save capture with ARP Requests for successful ARP Replay.
        Overwrites previous file, if any exists.
        :param source_arp_cap_path: path to capture of ARP Requests
        """
        if not os.path.isfile(source_arp_cap_path):
            raise FileNotFoundError
        shutil.move(source_arp_cap_path, self.default_arp_cap_path)
        self.arp_cap_path = self.default_arp_cap_path

    def save_psk_file(self, source_psk_file_path):
        """
        Save PSK file containing hexadecimal cracked key for network.
        Overwrites previous file, if any exists.
        :param source_psk_file_path: path to PSK file
        """
        if not os.path.isfile(source_psk_file_path):
            raise FileNotFoundError
        shutil.move(source_psk_file_path, self.default_psk_path)
        self.psk_path = self.default_psk_path

    def save_wpa_handshake_cap(self, source_wpa_handshake_cap_path):
        """
        Save capture with WPA handshake.
        Overwrites previous file, if any exists.
        :param source_wpa_handshake_cap_path: path to capture with WPA handshake
        """
        if not os.path.isfile(source_wpa_handshake_cap_path):
            raise FileNotFoundError
        shutil.move(source_wpa_handshake_cap_path, self.default_wpa_handshake_cap_path)
        self.wpa_handshake_cap_path = self.default_wpa_handshake_cap_path

    def save_prga_xor(self, source_prga_xor_path):
        """
        Save file containing PRGA XOR keystream.
        Overwrites previous file, if any exists.
        :param source_prga_xor_path: path to file containing PRGA XOR keystream
        """
        if not os.path.isfile(source_prga_xor_path):
            raise FileNotFoundError
        shutil.move(source_prga_xor_path, self.default_prga_xor_path)
        self.prga_xor_path = self.default_prga_xor_path

    def add_associated_station(self, station):
        station.associated_ap = self
        self.associated_stations.append(station)

    def update_known(self):
        """
        Update known network based on saved files.
        """
        if not self.arp_cap_path and os.path.isfile(self.default_arp_cap_path):
            self.arp_cap_path = self.default_arp_cap_path
            logger.debug(self.essid + ' arp_cap known')

        if not self.psk_path and os.path.isfile(self.default_psk_path):
            self.psk_path = self.default_psk_path
            logger.debug(self.essid + ' psk known')

        if not self.prga_xor_path and os.path.isfile(self.default_prga_xor_path):
            self.prga_xor_path = self.default_prga_xor_path
            logger.debug(self.essid + ' prga_xor known')

        if not self.wpa_handshake_cap_path and os.path.isfile(self.default_wpa_handshake_cap_path):
            self.wpa_handshake_cap_path = self.default_wpa_handshake_cap_path
            logger.debug(self.essid + ' wpa_handshake_cap known')

def interface_exists(name: str) -> bool:
    """
    Check if interface with given name exists.
    Does not check whether given name is *wireless* interface.
    :type name: str
    :param name: interface name
    :rtype: bool
    :return: True if interface exists, False otherwise.
    """
    return name in netifaces.interfaces()

class WirelessInterface(object):
    def __str__(self, *args, **kwargs):  # TODO (xvondr20) just for debugging
        s = 'WirelessInterface(' + ', '.join([
            self.name,
            self.mac_address,
            self.channel,
            self.driver,
            self.chipset
        ])
        if self.monitor_mode:
            s += ', monitor'
        s += ')'
        return s

    def __init__(self, name, driver=None, chipset=None):
        """
        Raises:
            ValueError if name is not a valid interface name
        """
        if not interface_exists(name):
            raise ValueError('You must specify a valid interface name.')

        self.name_original = name
        self.name_monitor = None

        # get MAC address
        self.mac_address_original = self.get_mac_by_name(self.name)
        self.mac_address_spoofed = None

        self.channel = None
        self.monitor_mode = False

        # additional data
        self.driver = driver
        self.chipset = chipset

    @staticmethod
    def get_wireless_interface_obj(interface):
        """
        Get WirelessInterface object based on provided argument.
        If interface is already WirelessInterface object, it is just returned. If interface is a valid interface name,
        appropriate WirelessInterface object is created and returned.
        :param interface: WirelessInterface object or string representing valid wireless interface name
        :return: WirelessInterface object
        Raises:
            ValueError if provided interface string is not a valid interface name
            TypeError if provided interface is not a string nor a WirelessInterface object
        """
        if isinstance(interface, WirelessInterface):
            return interface
        elif isinstance(interface, str):
            return WirelessInterface(name=interface)
        else:
            raise TypeError

    @staticmethod
    def get_mac_by_name(name):
        """
        Get MAC address of interface specified by name of the interface.
        :param name: name of the network interface
        :return: string MAC address
        """
        # TODO(xvondr20) Is this safe?
        ifa = netifaces.ifaddresses(name)
        mac = ifa[netifaces.AF_LINK][0]['addr']
        return mac

    @property
    def mac_address(self):
        """
        Get current MAC address.
        """
        assert self.mac_address_spoofed or self.mac_address_original, 'No MAC address available.'

        if self.mac_address_spoofed:
            return self.mac_address_spoofed
        else:
            return self.mac_address_original

    @property
    def name(self):
        """
        Get current interface name.
        """
        assert self.name_monitor or self.name_original, 'No interface name available.'

        if self.name_monitor:
            return self.name_monitor
        else:
            return self.name_original

    @property
    def gateway(self):
        """
        Get current default gateway.
        """
        gateway = None
        gws = netifaces.gateways()
        for gw in gws[netifaces.AF_INET]:
            if gw[1] == self.name:
                gateway = gw[0]
                break
        assert gateway, 'No default gateway available.'
        return gateway

    def start_monitor_mode(self, channel=None):
        """
        :param channel: monitor interface channel
        Raises:
            CalledProcessError if airmon-ng returncode is non-zero
        """
        assert not self.monitor_mode, 'Interface already in monitor mode.'

        cre_mon_enabled = re.compile(
            r'^\s+\(\S+ monitor mode vif enabled for \[\S+\](?P<name>\S+) on \[\S+\](?P<mon>\S+)\)$')

        cmd = ['airmon-ng', 'start', self.name]
        if channel:
            cmd.append(str(channel))
            self.channel = channel
        process = subprocess.run(cmd,
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                 universal_newlines=True)

        process.check_returncode()
        # check stderr
        # TODO (xvondr20) Does 'airmon-ng' ever print anything to stderr?
        assert process.stderr == ''

        for line in process.stdout.splitlines():
            m = cre_mon_enabled.match(line)
            if m:
                self.monitor_mode = True
                self.name_monitor = m.group('mon')
                break

    def stop_monitor_mode(self):
        """
        Raises:
            CalledProcessError if airmon-ng returncode is non-zero
        """
        assert self.monitor_mode, 'Interface not in monitor mode.'

        cre_mon_disabled = re.compile(r'^\s+\(\S+ monitor mode vif disabled for \[\S+\](?P<mon>\S+)\)$')

        cmd = ['airmon-ng', 'stop', self.name]
        process = subprocess.run(cmd,
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                 universal_newlines=True)

        process.check_returncode()
        # check stderr
        # TODO (xvondr20) Does 'airmon-ng' ever print anything to stderr?
        assert process.stderr == ''

        for line in process.stdout.splitlines():
            m = cre_mon_disabled.match(line)
            if m:
                self.monitor_mode = False
                self.name_monitor = None
                break
<|endoftext|>


<commit_before>import asyncio
import logging
import time
from asyncio import iscoroutinefunction, iscoroutine
from typing import List

from lightbus.api import Api
from lightbus.client import commands
from lightbus.client.commands import ConsumeRpcsCommand
from lightbus.client.subclients.base import BaseSubClient
from lightbus.client.utilities import validate_event_or_rpc_name, Error
from lightbus.client.validator import validate_outgoing, validate_incoming
from lightbus.exceptions import (
    NoApisToListenOn,
    SuddenDeathException,
    LightbusServerError,
    LightbusTimeout,
)
from lightbus.log import L, Bold
from lightbus.message import ResultMessage, RpcMessage
from lightbus.utilities.async_tools import run_user_provided_callable, cancel
from lightbus.utilities.casting import cast_to_signature
from lightbus.utilities.deforming import deform_to_bus
from lightbus.utilities.frozendict import frozendict
from lightbus.utilities.human import human_time
from lightbus.utilities.singledispatch import singledispatchmethod

logger = logging.getLogger(__name__)

async def bail_on_error(error_queue: asyncio.Queue, co):

    assert iscoroutine(co), "@bail_on_error only operates on coroutines"
    # TODO: May need to pass a fresh error queue through to ensure
    #       we don't compete with the bus client's event monitor
    fn_task = asyncio.ensure_future(co)
    monitor_task = asyncio.ensure_future(error_queue.get())

    done, pending = await asyncio.wait({fn_task, monitor_task}, return_when=asyncio.FIRST_COMPLETED)

    if fn_task in done:
        # All ok
        await cancel(monitor_task)
        return fn_task.result()
    else:
        # An error appeared in the queue
        await cancel(fn_task)
        error: Error = monitor_task.result()
        raise error.value

class RpcResultClient(BaseSubClient):
    """Functionality for both RPCs and results"""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    async def consume_rpcs(self, apis: List[Api] = None):
        """Start a background task to consume RPCs

        This will consumer RPCs on APIs which have been registered with this
        bus client.
        """
        if apis is None:
            apis = self.api_registry.all()

        if not apis:
            raise NoApisToListenOn(
                "No APIs to consume on in consume_rpcs(). Either this method was called with apis=[], "
                "or the API registry is empty."
            )

        api_names = [api.meta.name for api in apis]

        await self.producer.send(ConsumeRpcsCommand(api_names=api_names)).wait()

    async def call_rpc_remote(
        self, api_name: str, name: str, kwargs: dict = frozendict(), options: dict = frozendict()
    ):
        """ Perform an RPC call

        Call an RPC and return the result.
        """
        # TODO: InternalProducer command
        kwargs = deform_to_bus(kwargs)
        rpc_message = RpcMessage(api_name=api_name, procedure_name=name, kwargs=kwargs)
        validate_event_or_rpc_name(api_name, "rpc", name)

        logger.info("📞  Calling remote RPC {}.{}".format(Bold(api_name), Bold(name)))

        start_time = time.time()
        # TODO: It is possible that the RPC will be called before we start waiting for the
        #       response. This is bad.

        validate_outgoing(self.config, self.schema, rpc_message)

        await self._execute_hook("before_rpc_call", rpc_message=rpc_message)

        result_queue = asyncio.Queue()

        # Send the RPC
        await self.producer.send(
            commands.CallRpcCommand(message=rpc_message, options=options)
        ).wait()

        # Start a listener which will wait for results
        await self.producer.send(
            commands.ReceiveResultCommand(
                message=rpc_message, destination_queue=result_queue, options=options
            )
        ).wait()

        # Wait for the result from the listener we started.
        # The RpcResultDock will handle timeouts
        result = await bail_on_error(self.error_queue, result_queue.get())

        call_time = time.time() - start_time

        try:
            if isinstance(result, Exception):
                raise result
        except asyncio.TimeoutError:
            # TODO: Remove RPC from queue. Perhaps add a RpcBackend.cancel() method. Optional,
            #       as not all backends will support it. No point processing calls which have timed out.
            raise LightbusTimeout(
                f"Timeout when calling RPC {rpc_message.canonical_name} after waiting for {human_time(call_time)}. "
                f"It is possible no Lightbus process is serving this API, or perhaps it is taking "
                f"too long to process the request. In which case consider raising the 'rpc_timeout' "
                f"config option."
            ) from None
        else:
            assert isinstance(result, ResultMessage)
            result_message = result

        await self._execute_hook(
            "after_rpc_call", rpc_message=rpc_message, result_message=result_message
        )

        if not result_message.error:
            logger.info(
                L(
                    "🏁  Remote call of {} completed in {}",
                    Bold(rpc_message.canonical_name),
                    human_time(call_time),
                )
            )
        else:
            logger.warning(
                L(
                    "⚡ Server error during remote call of {}. Took {}: {}",
                    Bold(rpc_message.canonical_name),
                    human_time(call_time),
                    result_message.result,
                )
            )
            raise LightbusServerError(
                "Error while calling {}: {}\nRemote stack trace:\n{}".format(
                    rpc_message.canonical_name, result_message.result, result_message.trace
                )
            )

        validate_incoming(self.config, self.schema, result_message)

        return result_message.result

    async def _call_rpc_local(self, api_name: str, name: str, kwargs: dict = frozendict()):
        # TODO: InternalProducer command
        api = self.api_registry.get(api_name)
        validate_event_or_rpc_name(api_name, "rpc", name)

        start_time = time.time()
        try:
            method = getattr(api, name)
            if self.config.api(api_name).cast_values:
                kwargs = cast_to_signature(kwargs, method)
            result = await run_user_provided_callable(method, args=[], kwargs=kwargs)
        except (asyncio.CancelledError, SuddenDeathException):
            raise
        except Exception as e:
            logging.exception(e)
            logger.warning(
                L(
                    "⚡  Error while executing {}.{}. Took {}",
                    Bold(api_name),
                    Bold(name),
                    human_time(time.time() - start_time),
                )
            )
            raise
        else:
            logger.info(
                L(
                    "⚡  Executed {}.{} in {}",
                    Bold(api_name),
                    Bold(name),
                    human_time(time.time() - start_time),
                )
            )
            return result

    async def close(self):
        await self.producer.send(commands.CloseCommand()).wait()

        await self.consumer.close()
        await self.producer.close()

    @singledispatchmethod
    async def handle(self, command):
        raise NotImplementedError(f"Did not recognise command {command.__class__.__name__}")

    @handle.register
    async def handle_execute_rpc(self, command: commands.ExecuteRpcCommand):
        await self.schema.ensure_loaded_from_bus()
        validate_incoming(self.config, self.schema, command.message)

        await self._execute_hook("before_rpc_execution", rpc_message=command.message)
        try:
            result = await self._call_rpc_local(
                api_name=command.message.api_name,
                name=command.message.procedure_name,
                kwargs=command.message.kwargs,
            )
        except SuddenDeathException:
            # Used to simulate message failure for testing
            return
        except asyncio.CancelledError:
            raise
        except Exception as e:
            result = e
        else:
            result = deform_to_bus(result)

        result_message = ResultMessage(
            result=result,
            rpc_message_id=command.message.id,
            api_name=command.message.api_name,
            procedure_name=command.message.procedure_name,
        )
        await self._execute_hook(
            "after_rpc_execution", rpc_message=command.message, result_message=result_message
        )

        if not result_message.error:
            validate_outgoing(self.config, self.schema, result_message)

        await self.producer.send(
            commands.SendResultCommand(message=result_message, rpc_message=command.message)
        ).wait()
<commit_msg>Updating hook execution to use the new hook registry<commit_after>import asyncio
import logging
import time
from asyncio import iscoroutinefunction, iscoroutine
from typing import List

from lightbus.api import Api
from lightbus.client import commands
from lightbus.client.commands import ConsumeRpcsCommand
from lightbus.client.subclients.base import BaseSubClient
from lightbus.client.utilities import validate_event_or_rpc_name, Error
from lightbus.client.validator import validate_outgoing, validate_incoming
from lightbus.exceptions import (
    NoApisToListenOn,
    SuddenDeathException,
    LightbusServerError,
    LightbusTimeout,
)
from lightbus.log import L, Bold
from lightbus.message import ResultMessage, RpcMessage
from lightbus.utilities.async_tools import run_user_provided_callable, cancel
from lightbus.utilities.casting import cast_to_signature
from lightbus.utilities.deforming import deform_to_bus
from lightbus.utilities.frozendict import frozendict
from lightbus.utilities.human import human_time
from lightbus.utilities.singledispatch import singledispatchmethod

logger = logging.getLogger(__name__)

async def bail_on_error(error_queue: asyncio.Queue, co):

    assert iscoroutine(co), "@bail_on_error only operates on coroutines"
    # TODO: May need to pass a fresh error queue through to ensure
    #       we don't compete with the bus client's event monitor
    fn_task = asyncio.ensure_future(co)
    monitor_task = asyncio.ensure_future(error_queue.get())

    done, pending = await asyncio.wait({fn_task, monitor_task}, return_when=asyncio.FIRST_COMPLETED)

    if fn_task in done:
        # All ok
        await cancel(monitor_task)
        return fn_task.result()
    else:
        # An error appeared in the queue
        await cancel(fn_task)
        error: Error = monitor_task.result()
        raise error.value

class RpcResultClient(BaseSubClient):
    """Functionality for both RPCs and results"""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    async def consume_rpcs(self, apis: List[Api] = None):
        """Start a background task to consume RPCs

        This will consumer RPCs on APIs which have been registered with this
        bus client.
        """
        if apis is None:
            apis = self.api_registry.all()

        if not apis:
            raise NoApisToListenOn(
                "No APIs to consume on in consume_rpcs(). Either this method was called with apis=[], "
                "or the API registry is empty."
            )

        api_names = [api.meta.name for api in apis]

        await self.producer.send(ConsumeRpcsCommand(api_names=api_names)).wait()

    async def call_rpc_remote(
        self, api_name: str, name: str, kwargs: dict = frozendict(), options: dict = frozendict()
    ):
        """ Perform an RPC call

        Call an RPC and return the result.
        """
        # TODO: InternalProducer command
        kwargs = deform_to_bus(kwargs)
        rpc_message = RpcMessage(api_name=api_name, procedure_name=name, kwargs=kwargs)
        validate_event_or_rpc_name(api_name, "rpc", name)

        logger.info("📞  Calling remote RPC {}.{}".format(Bold(api_name), Bold(name)))

        start_time = time.time()
        # TODO: It is possible that the RPC will be called before we start waiting for the
        #       response. This is bad.

        validate_outgoing(self.config, self.schema, rpc_message)

        await self.hook_registry.execute("before_rpc_call", rpc_message=rpc_message)

        result_queue = asyncio.Queue()

        # Send the RPC
        await self.producer.send(
            commands.CallRpcCommand(message=rpc_message, options=options)
        ).wait()

        # Start a listener which will wait for results
        await self.producer.send(
            commands.ReceiveResultCommand(
                message=rpc_message, destination_queue=result_queue, options=options
            )
        ).wait()

        # Wait for the result from the listener we started.
        # The RpcResultDock will handle timeouts
        result = await bail_on_error(self.error_queue, result_queue.get())

        call_time = time.time() - start_time

        try:
            if isinstance(result, Exception):
                raise result
        except asyncio.TimeoutError:
            # TODO: Remove RPC from queue. Perhaps add a RpcBackend.cancel() method. Optional,
            #       as not all backends will support it. No point processing calls which have timed out.
            raise LightbusTimeout(
                f"Timeout when calling RPC {rpc_message.canonical_name} after waiting for {human_time(call_time)}. "
                f"It is possible no Lightbus process is serving this API, or perhaps it is taking "
                f"too long to process the request. In which case consider raising the 'rpc_timeout' "
                f"config option."
            ) from None
        else:
            assert isinstance(result, ResultMessage)
            result_message = result

        await self.hook_registry.execute(
            "after_rpc_call", rpc_message=rpc_message, result_message=result_message
        )

        if not result_message.error:
            logger.info(
                L(
                    "🏁  Remote call of {} completed in {}",
                    Bold(rpc_message.canonical_name),
                    human_time(call_time),
                )
            )
        else:
            logger.warning(
                L(
                    "⚡ Server error during remote call of {}. Took {}: {}",
                    Bold(rpc_message.canonical_name),
                    human_time(call_time),
                    result_message.result,
                )
            )
            raise LightbusServerError(
                "Error while calling {}: {}\nRemote stack trace:\n{}".format(
                    rpc_message.canonical_name, result_message.result, result_message.trace
                )
            )

        validate_incoming(self.config, self.schema, result_message)

        return result_message.result

    async def _call_rpc_local(self, api_name: str, name: str, kwargs: dict = frozendict()):
        # TODO: InternalProducer command
        api = self.api_registry.get(api_name)
        validate_event_or_rpc_name(api_name, "rpc", name)

        start_time = time.time()
        try:
            method = getattr(api, name)
            if self.config.api(api_name).cast_values:
                kwargs = cast_to_signature(kwargs, method)
            result = await run_user_provided_callable(method, args=[], kwargs=kwargs)
        except (asyncio.CancelledError, SuddenDeathException):
            raise
        except Exception as e:
            logging.exception(e)
            logger.warning(
                L(
                    "⚡  Error while executing {}.{}. Took {}",
                    Bold(api_name),
                    Bold(name),
                    human_time(time.time() - start_time),
                )
            )
            raise
        else:
            logger.info(
                L(
                    "⚡  Executed {}.{} in {}",
                    Bold(api_name),
                    Bold(name),
                    human_time(time.time() - start_time),
                )
            )
            return result

    async def close(self):
        await self.producer.send(commands.CloseCommand()).wait()

        await self.consumer.close()
        await self.producer.close()

    @singledispatchmethod
    async def handle(self, command):
        raise NotImplementedError(f"Did not recognise command {command.__class__.__name__}")

    @handle.register
    async def handle_execute_rpc(self, command: commands.ExecuteRpcCommand):
        await self.schema.ensure_loaded_from_bus()
        validate_incoming(self.config, self.schema, command.message)

        await self.hook_registry.execute("before_rpc_execution", rpc_message=command.message)
        try:
            result = await self._call_rpc_local(
                api_name=command.message.api_name,
                name=command.message.procedure_name,
                kwargs=command.message.kwargs,
            )
        except SuddenDeathException:
            # Used to simulate message failure for testing
            return
        except asyncio.CancelledError:
            raise
        except Exception as e:
            result = e
        else:
            result = deform_to_bus(result)

        result_message = ResultMessage(
            result=result,
            rpc_message_id=command.message.id,
            api_name=command.message.api_name,
            procedure_name=command.message.procedure_name,
        )
        await self.hook_registry.execute(
            "after_rpc_execution", rpc_message=command.message, result_message=result_message
        )

        if not result_message.error:
            validate_outgoing(self.config, self.schema, result_message)

        await self.producer.send(
            commands.SendResultCommand(message=result_message, rpc_message=command.message)
        ).wait()
<|endoftext|>


<commit_before># coding=utf-8
# Copyright 2018 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Batch of environments inside the TensorFlow graph."""

# The code was based on Danijar Hafner's code from tf.agents:
# https://github.com/tensorflow/agents/blob/master/agents/tools/in_graph_batch_env.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports

from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
from tensor2tensor.utils import registry
from tensor2tensor.utils import trainer_lib

import tensorflow as tf

flags = tf.flags
FLAGS = flags.FLAGS

class HistoryBuffer():

  def __init__(self, initial_frames):
    self._history_buff = None
    self.initial_frames = initial_frames
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      self._history_buff = tf.get_variable(
            "history_observ", initializer=self.initial_frames, trainable=False)

  def get_all_elements(self):
    return self._history_buff.read_value()

  def move_by_one_element(self, element):
    last_removed = self._history_buff.read_value()[1:, ...]
    element = tf.expand_dims(element, dim=0)
    moved = tf.concat([last_removed, element], axis=0)
    with tf.control_dependencies([moved]):
      with tf.control_dependencies([self._history_buff.assign(moved)]):
        return self._history_buff.read_value()

  def reset(self):
    with tf.control_dependencies([self._history_buff.assign(self.initial_frames)]):
      return self._history_buff.read_value()

class SimulatedBatchEnv(InGraphBatchEnv):
  """Batch of environments inside the TensorFlow graph.

  The batch of environments will be stepped and reset inside of the graph using
  a tf.py_func(). The current batch of observations, actions, rewards, and done
  flags are held in according variables.
  """
  NUMBER_OF_HISTORY_FRAMES = 4

  def __init__(self, environment_lambda, length):
    """Batch of environments inside the TensorFlow graph."""
    self.length = length

    #TODO(piotrmilos): For the moment we are fine with that.
    assert self.length==1, "Currently SimulatedBatchEnv support only one env"
    initialization_env = environment_lambda()
    hparams = trainer_lib.create_hparams(
        FLAGS.hparams_set, problem_name=FLAGS.problem, data_dir="UNUSED")
    hparams.force_full_predict = True
    self._model = registry.model(FLAGS.model)(
        hparams, tf.estimator.ModeKeys.PREDICT)

    self.action_space = initialization_env.action_space
    self.action_shape = list(initialization_env.action_space.shape)
    self.action_dtype = tf.int32

    obs = []
    if hasattr(initialization_env.env, "get_starting_data"):
      obs, _, _ = initialization_env.env.get_starting_data()
    else:
      # piotrmilos
      # Ancient method for environments not supporting get_starting_data
      # This is probably not compatibile with NUMBER_OF_HISTORY_FRAMES!=2
      # Should be removed at some point
      num_frames = SimulatedBatchEnv.NUMBER_OF_HISTORY_FRAMES
      initialization_env.reset()
      skip_frames = 20
      for _ in range(skip_frames):
        initialization_env.step(0)
      for i in range(num_frames):
        obs.append(initialization_env.step(0)[0])

    initial_frames = tf.stack(obs)
    initial_frames = tf.cast(initial_frames, tf.float32)

    self.history_buffer = HistoryBuffer(initial_frames)

    shape = (self.length,) + initialization_env.observation_space.shape
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      self._observ = tf.get_variable(
        "observ", shape, initializer=tf.zeros_initializer, trainable=False)

  def __len__(self):
    """Number of combined environments."""
    return self.length

  def simulate(self, action):
    with tf.name_scope("environment/simulate"):
      action = tf.expand_dims(action, axis=0)
      actions = [action]*SimulatedBatchEnv.NUMBER_OF_HISTORY_FRAMES
      actions = tf.concat(actions, axis=0)
      history = self.history_buffer.get_all_elements()
      inputs = {"inputs": tf.expand_dims(history, axis=0),  # Add batch.
                "input_action": tf.expand_dims(actions, axis=0)}
      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
        model_output = self._model.infer(inputs)
      observ = model_output["targets"]
      observ = tf.cast(observ[:, 0, :, :, :], tf.float32)
      # TODO(lukaszkaiser): instead of -1 use min_reward in the line below.
      reward = model_output["target_reward"][:, 0, 0, 0] - 1
      reward = tf.cast(reward, tf.float32)
      # Some wrappers need explicit shape, so we reshape here.
      reward = tf.reshape(reward, shape=(self.length,))
      done = tf.constant(False, tf.bool, shape=(self.length,))

      with tf.control_dependencies([observ]):
          with tf.control_dependencies([self._observ.assign(observ),
                                        self.history_buffer.move_by_one_element(observ[0, ...])]):
            return tf.identity(reward), tf.identity(done)

  def reset(self, indices=None):
    """Reset the batch of environments.

    Args:
      indices: The batch indices of the environments to reset.

    Returns:
      Batch tensor of the new observations.
    """
    return tf.cond(
        tf.cast(tf.shape(indices)[0], tf.bool),
        lambda: self._reset_non_empty(indices), lambda: 0.0)

  def _reset_non_empty(self, indices):
    """Reset the batch of environments.

    Args:
      indices: The batch indices of the environments to reset; defaults to all.

    Returns:
      Batch tensor of the new observations.
    """
    with tf.control_dependencies([self.history_buffer.reset()]):
      with tf.control_dependencies([self._observ.assign(0.0*self._observ)]):
        ooo = tf.Print(self._observ.read_value(), [1], message="reset op")
        return tf.identity(ooo)

  @property
  def observ(self):
    """Access the variable holding the current observation."""
    return self._observ
<commit_msg>removing stupid log<commit_after># coding=utf-8
# Copyright 2018 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Batch of environments inside the TensorFlow graph."""

# The code was based on Danijar Hafner's code from tf.agents:
# https://github.com/tensorflow/agents/blob/master/agents/tools/in_graph_batch_env.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports

from tensor2tensor.rl.envs.in_graph_batch_env import InGraphBatchEnv
from tensor2tensor.utils import registry
from tensor2tensor.utils import trainer_lib

import tensorflow as tf

flags = tf.flags
FLAGS = flags.FLAGS

class HistoryBuffer():

  def __init__(self, initial_frames):
    self._history_buff = None
    self.initial_frames = initial_frames
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      self._history_buff = tf.get_variable(
            "history_observ", initializer=self.initial_frames, trainable=False)

  def get_all_elements(self):
    return self._history_buff.read_value()

  def move_by_one_element(self, element):
    last_removed = self._history_buff.read_value()[1:, ...]
    element = tf.expand_dims(element, dim=0)
    moved = tf.concat([last_removed, element], axis=0)
    with tf.control_dependencies([moved]):
      with tf.control_dependencies([self._history_buff.assign(moved)]):
        return self._history_buff.read_value()

  def reset(self):
    with tf.control_dependencies([self._history_buff.assign(self.initial_frames)]):
      return self._history_buff.read_value()

class SimulatedBatchEnv(InGraphBatchEnv):
  """Batch of environments inside the TensorFlow graph.

  The batch of environments will be stepped and reset inside of the graph using
  a tf.py_func(). The current batch of observations, actions, rewards, and done
  flags are held in according variables.
  """
  NUMBER_OF_HISTORY_FRAMES = 4

  def __init__(self, environment_lambda, length):
    """Batch of environments inside the TensorFlow graph."""
    self.length = length

    #TODO(piotrmilos): For the moment we are fine with that.
    assert self.length==1, "Currently SimulatedBatchEnv support only one env"
    initialization_env = environment_lambda()
    hparams = trainer_lib.create_hparams(
        FLAGS.hparams_set, problem_name=FLAGS.problem, data_dir="UNUSED")
    hparams.force_full_predict = True
    self._model = registry.model(FLAGS.model)(
        hparams, tf.estimator.ModeKeys.PREDICT)

    self.action_space = initialization_env.action_space
    self.action_shape = list(initialization_env.action_space.shape)
    self.action_dtype = tf.int32

    obs = []
    if hasattr(initialization_env.env, "get_starting_data"):
      obs, _, _ = initialization_env.env.get_starting_data()
    else:
      # piotrmilos
      # Ancient method for environments not supporting get_starting_data
      # This is probably not compatibile with NUMBER_OF_HISTORY_FRAMES!=2
      # Should be removed at some point
      num_frames = SimulatedBatchEnv.NUMBER_OF_HISTORY_FRAMES
      initialization_env.reset()
      skip_frames = 20
      for _ in range(skip_frames):
        initialization_env.step(0)
      for i in range(num_frames):
        obs.append(initialization_env.step(0)[0])

    initial_frames = tf.stack(obs)
    initial_frames = tf.cast(initial_frames, tf.float32)

    self.history_buffer = HistoryBuffer(initial_frames)

    shape = (self.length,) + initialization_env.observation_space.shape
    with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
      self._observ = tf.get_variable(
        "observ", shape, initializer=tf.zeros_initializer, trainable=False)

  def __len__(self):
    """Number of combined environments."""
    return self.length

  def simulate(self, action):
    with tf.name_scope("environment/simulate"):
      action = tf.expand_dims(action, axis=0)
      actions = [action]*SimulatedBatchEnv.NUMBER_OF_HISTORY_FRAMES
      actions = tf.concat(actions, axis=0)
      history = self.history_buffer.get_all_elements()
      inputs = {"inputs": tf.expand_dims(history, axis=0),  # Add batch.
                "input_action": tf.expand_dims(actions, axis=0)}
      with tf.variable_scope(tf.get_variable_scope(), reuse=tf.AUTO_REUSE):
        model_output = self._model.infer(inputs)
      observ = model_output["targets"]
      observ = tf.cast(observ[:, 0, :, :, :], tf.float32)
      # TODO(lukaszkaiser): instead of -1 use min_reward in the line below.
      reward = model_output["target_reward"][:, 0, 0, 0] - 1
      reward = tf.cast(reward, tf.float32)
      # Some wrappers need explicit shape, so we reshape here.
      reward = tf.reshape(reward, shape=(self.length,))
      done = tf.constant(False, tf.bool, shape=(self.length,))

      with tf.control_dependencies([observ]):
          with tf.control_dependencies([self._observ.assign(observ),
                                        self.history_buffer.move_by_one_element(observ[0, ...])]):
            return tf.identity(reward), tf.identity(done)

  def reset(self, indices=None):
    """Reset the batch of environments.

    Args:
      indices: The batch indices of the environments to reset.

    Returns:
      Batch tensor of the new observations.
    """
    return tf.cond(
        tf.cast(tf.shape(indices)[0], tf.bool),
        lambda: self._reset_non_empty(indices), lambda: 0.0)

  def _reset_non_empty(self, indices):
    """Reset the batch of environments.

    Args:
      indices: The batch indices of the environments to reset; defaults to all.

    Returns:
      Batch tensor of the new observations.
    """
    with tf.control_dependencies([self.history_buffer.reset()]):
      with tf.control_dependencies([self._observ.assign(0.0*self._observ)]):
        return tf.identity(self._observ.read_value())

  @property
  def observ(self):
    """Access the variable holding the current observation."""
    return self._observ
<|endoftext|>


<commit_before><commit_msg>Handle unattached nodes; default to Spoke permissions<commit_after><|endoftext|>