<commit_before>import fontTools.ttLib as ttLib
from fontTools.ttLib.ttVisitor import TTVisitor
from fontTools.misc.textTools import Tag

class JsonVisitor(TTVisitor):

    def _open(self, s):
        print(s, file=self.file)
        self._indent += self.indent

    def _close(self, s):
        self._indent = self._indent[:-len(self.indent)]
        print("%s%s," % (self._indent, s), file=self.file)

    def __init__(self, file, indent="  "):
        self.file = file
        self.indent = indent
        self._indent = ""

    def visitObject(self, obj):
        self._open("{")
        super().visitObject(obj)
        print('%s"type": "%s"' % (self._indent, obj.__class__.__name__), file=self.file)
        self._close("}")

    def visitAttr(self, obj, attr, value):
        print('%s"%s": ' % (self._indent, attr), end="", file=self.file)
        self.visit(value)

    def visitList(self, obj, *args, **kwargs):
        self._open("[")
        for value in obj:
            print(self._indent, end="", file=self.file)
            self.visit(value, *args, **kwargs)
        self._close("]")

    def visitDict(self, obj, *args, **kwargs):
        self._open("{")
        for key, value in obj.items():
            print('%s"%s": ' % (self._indent, key), end="", file=self.file)
            self.visit(value, *args, **kwargs)
        self._close("}")

    def visitLeaf(self, obj):
        if isinstance(obj, tuple):
            obj = list(obj)

        if isinstance(obj, bytes):
            s = repr(obj)
            s = s[1:]
        else:
            s = repr(obj)

        if s[0] == "'":
            s = '"' + s[1:-1] + '"'

        print("%s," % s, file=self.file)

@JsonVisitor.register(ttLib.TTFont)
def visit(self, font):
    if hasattr(visitor, "font"):
        return False
    visitor.font = font

    self._open("{")
    for tag in font.keys():
        print('%s"%s": ' % (self._indent, tag), end="", file=self.file)
        visitor.visit(font[tag])
    self._close("}")

    del visitor.font
    return False

@JsonVisitor.register_attr(ttLib.getTableClass("glyf"), "glyphOrder")
def visit(visitor, obj, attr, value):
    return False

@JsonVisitor.register(ttLib.getTableModule("glyf").GlyphCoordinates)
def visit(self, obj):
    self.visitList(obj)
    return False

@JsonVisitor.register(Tag)
def visit(self, obj):
    print('"%s",' % str(obj), file=self.file)
    return False

if __name__ == "__main__":

    from fontTools.ttLib import TTFont
    import sys

    if len(sys.argv) != 2:
        print("usage: print-json.py font")
        sys.exit()

    font = TTFont(sys.argv[1])

    visitor = JsonVisitor(sys.stdout)
    visitor.visit(font)
<commit_msg>[print-json] Towards valid json syntax<commit_after>import fontTools.ttLib as ttLib
from fontTools.ttLib.ttVisitor import TTVisitor
from fontTools.misc.textTools import Tag

class JsonVisitor(TTVisitor):
    def _open(self, s):
        print(s, file=self.file)
        self._indent += self.indent
        self.comma = False

    def _close(self, s):
        self._indent = self._indent[: -len(self.indent)]
        print("\n%s%s" % (self._indent, s), end="", file=self.file)
        self.comma = True

    def __init__(self, file, indent="  "):
        self.file = file
        self.indent = indent
        self._indent = ""

    def visitObject(self, obj):
        self._open("{")
        super().visitObject(obj)
        if self.comma:
            print(",", end="", file=self.file)
        print(
            '\n%s"type": "%s"' % (self._indent, obj.__class__.__name__),
            end="",
            file=self.file,
        )
        self._close("}")

    def visitAttr(self, obj, attr, value):
        if self.comma:
            print(",", file=self.file)
        print('%s"%s": ' % (self._indent, attr), end="", file=self.file)
        self.visit(value)
        self.comma = True

    def visitList(self, obj, *args, **kwargs):
        self._open("[")
        comma = False
        for value in obj:
            if comma:
                print(",", end="", file=self.file)
                print(file=self.file)
            print(self._indent, end="", file=self.file)
            self.visit(value, *args, **kwargs)
            comma = True
        self._close("]")

    def visitDict(self, obj, *args, **kwargs):
        self._open("{")
        comma = False
        for key, value in obj.items():
            if comma:
                print(",", end="", file=self.file)
                print(file=self.file)
            print('%s"%s": ' % (self._indent, key), end="", file=self.file)
            self.visit(value, *args, **kwargs)
            comma = True
        self._close("}")

    def visitLeaf(self, obj):
        if isinstance(obj, tuple):
            obj = list(obj)

        if obj is None:
            s = "null"
        elif obj == True:
            s = "true"
        elif obj == False:
            s = "false"
        elif isinstance(obj, bytes):
            s = repr(obj)
            s = s[1:]
        else:
            s = repr(obj)

        if s[0] == "'":
            s = '"' + s[1:-1] + '"'

        print("%s" % s, end="", file=self.file)

@JsonVisitor.register(ttLib.TTFont)
def visit(self, font):
    if hasattr(visitor, "font"):
        print("{}", end="", file=self.file)
        return False
    visitor.font = font

    self._open("{")
    for tag in font.keys():
        if self.comma:
            print(",", file=self.file)
        print('\n%s"%s": ' % (self._indent, tag), end="", file=self.file)
        visitor.visit(font[tag])
    self._close("}")

    del visitor.font
    return False

@JsonVisitor.register_attr(ttLib.getTableClass("glyf"), "glyphOrder")
def visit(visitor, obj, attr, value):
    return False

@JsonVisitor.register(ttLib.getTableModule("glyf").GlyphCoordinates)
def visit(self, obj):
    self.visitList(obj)
    return False

@JsonVisitor.register(Tag)
def visit(self, obj):
    print('"%s"' % str(obj), end="", file=self.file)
    return False

if __name__ == "__main__":

    from fontTools.ttLib import TTFont
    import sys

    if len(sys.argv) != 2:
        print("usage: print-json.py font")
        sys.exit()

    font = TTFont(sys.argv[1])

    visitor = JsonVisitor(sys.stdout)
    visitor.visit(font)
<|endoftext|>


<commit_before>#
# Copyright (C) 2015, 2016 Satoru SATOH <ssato redhat.com>
# License: MIT
#
"""anyconfig.schema module.

.. versionchanged:: 0.8.3
   Replace format_checker with cls option

.. versionchanged:: 0.7.0
   allow passing `ac_schema_strict` to API :func:`gen_schema` to generate more
   strict and precise JSON schema object

.. versionadded:: 0.0.11
   Added new API :func:`gen_schema` to generate schema object

.. versionadded:: 0.0.10
   Added new API :func:`validate` to validate config with JSON schema
"""
from __future__ import absolute_import
try:
    import jsonschema
except ImportError:
    pass

import anyconfig.compat

_SIMPLETYPE_MAP = {list: "array", tuple: "array",
                   bool: "boolean",
                   int: "integer", float: "number",
                   dict: "object",
                   str: "string"}
_SIMPLE_TYPES = (bool, int, float, str)

if not anyconfig.compat.IS_PYTHON_3:
    try:
        _SIMPLETYPE_MAP[unicode] = "string"
        _SIMPLE_TYPES = (bool, int, float, str, unicode)
    except NameError:
        pass

def validate(data, schema, cls=None, **options):
    """
    Validate target object with given schema object, loaded from JSON schema.

    See also: https://python-jsonschema.readthedocs.org/en/latest/validate/

    :parae data: Target object (a dict or a dict-like object) to validate
    :param schema: Schema object (a dict or a dict-like object)
        instantiated from schema JSON file or schema JSON string
    :param cls:
        The class inherits :class:`~jsonschema.IValidator` that will be used to
        validate the instance.

    :param options: Other keyword options such as:

        - safe: Exception (jsonschema.ValidationError or jsonschema.SchemaError
          or others) will be thrown during validation process due to any
          validation or related errors. However, these will be catched by
          default, and will be re-raised if `safe` is False.

    :return: (True if validation succeeded else False, error message)
    """
    if not isinstance(cls, jsonschema.IValidator):
        cls = None
    try:
        try:
            jsonschema.validate(data, schema, cls=cls, **options)
            return (True, '')
        except (jsonschema.ValidationError, jsonschema.SchemaError,
                Exception) as exc:
            if options.get("safe", True):
                return (False, str(exc))
            else:
                raise

    except NameError:
        return (True, "Validation module (jsonschema) is not available")

    return (True, '')

def _process_options(**options):
    """
    Helper function to process keyword arguments passed to gen_schema.

    :return: A tuple of (typemap :: dict, strict :: bool)
    """
    return (options.get("ac_schema_typemap", _SIMPLETYPE_MAP),
            bool(options.get("ac_schema_strict", False)))

def array_to_schema(arr, **options):
    """
    Generate a JSON schema object with type annotation added for given object.

    :param arr: Array of dict or MergeableDict objects
    :param options: Other keyword options such as:

        - ac_schema_strict: True if more strict (precise) schema is needed
        - ac_schema_typemap: Type to JSON schema type mappings

    :return: Another MergeableDict instance represents JSON schema of items
    """
    (typemap, strict) = _process_options(**options)

    arr = list(arr)
    scm = dict(type=typemap[list],
               items=gen_schema(arr[0] if arr else "str", **options))
    if strict:
        nitems = len(arr)
        scm["minItems"] = nitems
        scm["uniqueItems"] = len(set(arr)) == nitems

    return scm

def object_to_schema(obj, **options):
    """
    Generate a node represents JSON schema object with type annotation added
    for given object node.

    :param obj: Dict or MergeableDict object
    :param options: Other keyword options such as:

        - ac_schema_strict: True if more strict (precise) schema is needed
        - ac_schema_typemap: Type to JSON schema type mappings

    :yield: Another MergeableDict instance represents JSON schema of object
    """
    (typemap, strict) = _process_options(**options)

    props = dict((k, gen_schema(v, **options)) for k, v in obj.items())
    scm = dict(type=typemap[dict], properties=props)
    if strict:
        scm["required"] = sorted(props.keys())

    return scm

def gen_schema(data, **options):
    """
    Generate a node represents JSON schema object with type annotation added
    for given object node.

    :param data: Configuration data object (dict[-like] or namedtuple)
    :param options: Other keyword options such as:

        - ac_schema_strict: True if more strict (precise) schema is needed
        - ac_schema_typemap: Type to JSON schema type mappings

    :return: A dict represents JSON schema of this node
    """
    if data is None:
        return dict(type="null")

    _type = type(data)

    if _type in _SIMPLE_TYPES:
        typemap = options.get("ac_schema_typemap", _SIMPLETYPE_MAP)
        scm = dict(type=typemap[_type])

    elif isinstance(data, dict):
        scm = object_to_schema(data, **options)

    elif _type in (list, tuple) or hasattr(data, "__iter__"):
        scm = array_to_schema(data, **options)

    return scm

# vim:sw=4:ts=4:et:
<commit_msg>fix: jsonschema.IValidator is newly introduced and not yet available in jsonschema 2.5.x I'm testing<commit_after>#
# Copyright (C) 2015, 2016 Satoru SATOH <ssato redhat.com>
# License: MIT
#
"""anyconfig.schema module.

.. versionchanged:: 0.8.3
   Replace format_checker with cls option

.. versionchanged:: 0.7.0
   allow passing `ac_schema_strict` to API :func:`gen_schema` to generate more
   strict and precise JSON schema object

.. versionadded:: 0.0.11
   Added new API :func:`gen_schema` to generate schema object

.. versionadded:: 0.0.10
   Added new API :func:`validate` to validate config with JSON schema
"""
from __future__ import absolute_import
try:
    import jsonschema
except ImportError:
    pass

import anyconfig.compat

_SIMPLETYPE_MAP = {list: "array", tuple: "array",
                   bool: "boolean",
                   int: "integer", float: "number",
                   dict: "object",
                   str: "string"}
_SIMPLE_TYPES = (bool, int, float, str)

if not anyconfig.compat.IS_PYTHON_3:
    try:
        _SIMPLETYPE_MAP[unicode] = "string"
        _SIMPLE_TYPES = (bool, int, float, str, unicode)
    except NameError:
        pass

def validate(data, schema, **options):
    """
    Validate target object with given schema object, loaded from JSON schema.

    See also: https://python-jsonschema.readthedocs.org/en/latest/validate/

    :parae data: Target object (a dict or a dict-like object) to validate
    :param schema: Schema object (a dict or a dict-like object)
        instantiated from schema JSON file or schema JSON string
    :param options: Other keyword options such as:

        - safe: Exception (jsonschema.ValidationError or jsonschema.SchemaError
          or others) will be thrown during validation process due to any
          validation or related errors. However, these will be catched by
          default, and will be re-raised if `safe` is False.

    :return: (True if validation succeeded else False, error message)
    """
    try:
        try:
            jsonschema.validate(data, schema, **options)
            return (True, '')
        except (jsonschema.ValidationError, jsonschema.SchemaError,
                Exception) as exc:
            if options.get("safe", True):
                return (False, str(exc))
            else:
                raise

    except NameError:
        return (True, "Validation module (jsonschema) is not available")

    return (True, '')

def _process_options(**options):
    """
    Helper function to process keyword arguments passed to gen_schema.

    :return: A tuple of (typemap :: dict, strict :: bool)
    """
    return (options.get("ac_schema_typemap", _SIMPLETYPE_MAP),
            bool(options.get("ac_schema_strict", False)))

def array_to_schema(arr, **options):
    """
    Generate a JSON schema object with type annotation added for given object.

    :param arr: Array of dict or MergeableDict objects
    :param options: Other keyword options such as:

        - ac_schema_strict: True if more strict (precise) schema is needed
        - ac_schema_typemap: Type to JSON schema type mappings

    :return: Another MergeableDict instance represents JSON schema of items
    """
    (typemap, strict) = _process_options(**options)

    arr = list(arr)
    scm = dict(type=typemap[list],
               items=gen_schema(arr[0] if arr else "str", **options))
    if strict:
        nitems = len(arr)
        scm["minItems"] = nitems
        scm["uniqueItems"] = len(set(arr)) == nitems

    return scm

def object_to_schema(obj, **options):
    """
    Generate a node represents JSON schema object with type annotation added
    for given object node.

    :param obj: Dict or MergeableDict object
    :param options: Other keyword options such as:

        - ac_schema_strict: True if more strict (precise) schema is needed
        - ac_schema_typemap: Type to JSON schema type mappings

    :yield: Another MergeableDict instance represents JSON schema of object
    """
    (typemap, strict) = _process_options(**options)

    props = dict((k, gen_schema(v, **options)) for k, v in obj.items())
    scm = dict(type=typemap[dict], properties=props)
    if strict:
        scm["required"] = sorted(props.keys())

    return scm

def gen_schema(data, **options):
    """
    Generate a node represents JSON schema object with type annotation added
    for given object node.

    :param data: Configuration data object (dict[-like] or namedtuple)
    :param options: Other keyword options such as:

        - ac_schema_strict: True if more strict (precise) schema is needed
        - ac_schema_typemap: Type to JSON schema type mappings

    :return: A dict represents JSON schema of this node
    """
    if data is None:
        return dict(type="null")

    _type = type(data)

    if _type in _SIMPLE_TYPES:
        typemap = options.get("ac_schema_typemap", _SIMPLETYPE_MAP)
        scm = dict(type=typemap[_type])

    elif isinstance(data, dict):
        scm = object_to_schema(data, **options)

    elif _type in (list, tuple) or hasattr(data, "__iter__"):
        scm = array_to_schema(data, **options)

    return scm

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


<commit_before>import socket
from argparse import ArgumentTypeError as Fatal
from unittest.mock import patch

import pytest

import sshuttle.options

_ip4_reprs = {
        '0.0.0.0': '0.0.0.0',
        '255.255.255.255': '255.255.255.255',
        '10.0': '10.0.0.0',
        '184.172.10.74': '184.172.10.74',
        '3098282570': '184.172.10.74',
        '0xb8.0xac.0x0a.0x4a': '184.172.10.74',
        '0270.0254.0012.0112': '184.172.10.74',
}

_ip4_swidths = (1, 8, 22, 27, 32)

_ip6_reprs = {
        '::': '::',
        '::1': '::1',
        'fc00::': 'fc00::',
        '2a01:7e00:e000:188::1': '2a01:7e00:e000:188::1'
}

_ip6_swidths = (48, 64, 96, 115, 128)

def _mock_getaddrinfo(host, *_):
    return {
        "example.com": [
            (socket.AF_INET6, socket.SOCK_STREAM, 0, '', ('2606:2800:220:1:248:1893:25c8:1946', 0, 0, 0)),
            (socket.AF_INET, socket.SOCK_STREAM, 0, '', ('93.184.216.34', 0)),
        ],
    }.get(host, [])

def test_parse_subnetport_ip4():
    for ip_repr, ip in _ip4_reprs.items():
        assert sshuttle.options.parse_subnetport(ip_repr) \
                == [(socket.AF_INET, ip, 32, 0, 0)]
    with pytest.raises(Fatal) as excinfo:
        sshuttle.options.parse_subnetport('10.256.0.0')
    assert str(excinfo.value) == 'Unable to resolve address: 10.256.0.0'

def test_parse_subnetport_ip4_with_mask():
    for ip_repr, ip in _ip4_reprs.items():
        for swidth in _ip4_swidths:
            assert sshuttle.options.parse_subnetport(
                    '/'.join((ip_repr, str(swidth)))
                    ) == [(socket.AF_INET, ip, swidth, 0, 0)]
    assert sshuttle.options.parse_subnetport('0/0') \
        == [(socket.AF_INET, '0.0.0.0', 0, 0, 0)]
    with pytest.raises(Fatal) as excinfo:
        sshuttle.options.parse_subnetport('10.0.0.0/33')
    assert str(excinfo.value) \
        == 'Slash in CIDR notation (/33) is not between 0 and 32'

def test_parse_subnetport_ip4_with_port():
    for ip_repr, ip in _ip4_reprs.items():
        assert sshuttle.options.parse_subnetport(':'.join((ip_repr, '80'))) \
            == [(socket.AF_INET, ip, 32, 80, 80)]
        assert sshuttle.options.parse_subnetport(':'.join((ip_repr, '80-90')))\
            == [(socket.AF_INET, ip, 32, 80, 90)]

def test_parse_subnetport_ip4_with_mask_and_port():
    for ip_repr, ip in _ip4_reprs.items():
        assert sshuttle.options.parse_subnetport(ip_repr + '/32:80') \
            == [(socket.AF_INET, ip, 32, 80, 80)]
        assert sshuttle.options.parse_subnetport(ip_repr + '/16:80-90') \
            == [(socket.AF_INET, ip, 16, 80, 90)]

def test_parse_subnetport_ip6():
    for ip_repr, ip in _ip6_reprs.items():
        assert sshuttle.options.parse_subnetport(ip_repr) \
                == [(socket.AF_INET6, ip, 128, 0, 0)]

def test_parse_subnetport_ip6_with_mask():
    for ip_repr, ip in _ip6_reprs.items():
        for swidth in _ip4_swidths + _ip6_swidths:
            assert sshuttle.options.parse_subnetport(
                    '/'.join((ip_repr, str(swidth)))
                    ) == [(socket.AF_INET6, ip, swidth, 0, 0)]
    assert sshuttle.options.parse_subnetport('::/0') \
        == [(socket.AF_INET6, '::', 0, 0, 0)]
    with pytest.raises(Fatal) as excinfo:
        sshuttle.options.parse_subnetport('fc00::/129')
    assert str(excinfo.value) \
        == 'Slash in CIDR notation (/129) is not between 0 and 128'

def test_parse_subnetport_ip6_with_port():
    for ip_repr, ip in _ip6_reprs.items():
        assert sshuttle.options.parse_subnetport('[' + ip_repr + ']:80') \
            == [(socket.AF_INET6, ip, 128, 80, 80)]
        assert sshuttle.options.parse_subnetport('[' + ip_repr + ']:80-90') \
            == [(socket.AF_INET6, ip, 128, 80, 90)]

def test_parse_subnetport_ip6_with_mask_and_port():
    for ip_repr, ip in _ip6_reprs.items():
        assert sshuttle.options.parse_subnetport('[' + ip_repr + '/128]:80') \
            == [(socket.AF_INET6, ip, 128, 80, 80)]
        assert sshuttle.options.parse_subnetport('[' + ip_repr + '/16]:80-90')\
            == [(socket.AF_INET6, ip, 16, 80, 90)]

def test_convert_arg_line_to_args_skips_comments():
    parser = sshuttle.options.MyArgumentParser()
    assert parser.convert_arg_line_to_args("# whatever something") == []

@patch('sshuttle.options.socket.getaddrinfo', side_effect = _mock_getaddrinfo)
def test_parse_subnetport_host(mock_getaddrinfo):
    assert set(sshuttle.options.parse_subnetport('example.com')) \
        == set([
            (socket.AF_INET6, '2606:2800:220:1:248:1893:25c8:1946', 128, 0, 0),
            (socket.AF_INET, '93.184.216.34', 32, 0, 0),
        ])

@patch('sshuttle.options.socket.getaddrinfo', side_effect = _mock_getaddrinfo)
def test_parse_subnetport_host_with_port(mock_getaddrinfo):
    assert set(sshuttle.options.parse_subnetport('example.com:80')) \
        == set([
            (socket.AF_INET6, '2606:2800:220:1:248:1893:25c8:1946', 128, 80, 80),
            (socket.AF_INET, '93.184.216.34', 32, 80, 80),
        ])
    assert set(sshuttle.options.parse_subnetport('example.com:80-90')) \
        == set([
            (socket.AF_INET6, '2606:2800:220:1:248:1893:25c8:1946', 128, 80, 90),
            (socket.AF_INET, '93.184.216.34', 32, 80, 90),
        ])
<commit_msg>add another example for host resolution tests<commit_after>import socket
from argparse import ArgumentTypeError as Fatal
from unittest.mock import patch

import pytest

import sshuttle.options

_ip4_reprs = {
        '0.0.0.0': '0.0.0.0',
        '255.255.255.255': '255.255.255.255',
        '10.0': '10.0.0.0',
        '184.172.10.74': '184.172.10.74',
        '3098282570': '184.172.10.74',
        '0xb8.0xac.0x0a.0x4a': '184.172.10.74',
        '0270.0254.0012.0112': '184.172.10.74',
}

_ip4_swidths = (1, 8, 22, 27, 32)

_ip6_reprs = {
        '::': '::',
        '::1': '::1',
        'fc00::': 'fc00::',
        '2a01:7e00:e000:188::1': '2a01:7e00:e000:188::1'
}

_ip6_swidths = (48, 64, 96, 115, 128)

def _mock_getaddrinfo(host, *_):
    return {
        "example.com": [
            (socket.AF_INET6, socket.SOCK_STREAM, 0, '', ('2606:2800:220:1:248:1893:25c8:1946', 0, 0, 0)),
            (socket.AF_INET, socket.SOCK_STREAM, 0, '', ('93.184.216.34', 0)),
        ],
        "my.local": [
            (socket.AF_INET6, socket.SOCK_STREAM, 0, '', ('::1', 0, 0, 0)),
            (socket.AF_INET, socket.SOCK_STREAM, 0, '', ('127.0.0.1', 0)),
        ],
    }.get(host, [])

def test_parse_subnetport_ip4():
    for ip_repr, ip in _ip4_reprs.items():
        assert sshuttle.options.parse_subnetport(ip_repr) \
                == [(socket.AF_INET, ip, 32, 0, 0)]
    with pytest.raises(Fatal) as excinfo:
        sshuttle.options.parse_subnetport('10.256.0.0')
    assert str(excinfo.value) == 'Unable to resolve address: 10.256.0.0'

def test_parse_subnetport_ip4_with_mask():
    for ip_repr, ip in _ip4_reprs.items():
        for swidth in _ip4_swidths:
            assert sshuttle.options.parse_subnetport(
                    '/'.join((ip_repr, str(swidth)))
                    ) == [(socket.AF_INET, ip, swidth, 0, 0)]
    assert sshuttle.options.parse_subnetport('0/0') \
        == [(socket.AF_INET, '0.0.0.0', 0, 0, 0)]
    with pytest.raises(Fatal) as excinfo:
        sshuttle.options.parse_subnetport('10.0.0.0/33')
    assert str(excinfo.value) \
        == 'Slash in CIDR notation (/33) is not between 0 and 32'

def test_parse_subnetport_ip4_with_port():
    for ip_repr, ip in _ip4_reprs.items():
        assert sshuttle.options.parse_subnetport(':'.join((ip_repr, '80'))) \
            == [(socket.AF_INET, ip, 32, 80, 80)]
        assert sshuttle.options.parse_subnetport(':'.join((ip_repr, '80-90')))\
            == [(socket.AF_INET, ip, 32, 80, 90)]

def test_parse_subnetport_ip4_with_mask_and_port():
    for ip_repr, ip in _ip4_reprs.items():
        assert sshuttle.options.parse_subnetport(ip_repr + '/32:80') \
            == [(socket.AF_INET, ip, 32, 80, 80)]
        assert sshuttle.options.parse_subnetport(ip_repr + '/16:80-90') \
            == [(socket.AF_INET, ip, 16, 80, 90)]

def test_parse_subnetport_ip6():
    for ip_repr, ip in _ip6_reprs.items():
        assert sshuttle.options.parse_subnetport(ip_repr) \
                == [(socket.AF_INET6, ip, 128, 0, 0)]

def test_parse_subnetport_ip6_with_mask():
    for ip_repr, ip in _ip6_reprs.items():
        for swidth in _ip4_swidths + _ip6_swidths:
            assert sshuttle.options.parse_subnetport(
                    '/'.join((ip_repr, str(swidth)))
                    ) == [(socket.AF_INET6, ip, swidth, 0, 0)]
    assert sshuttle.options.parse_subnetport('::/0') \
        == [(socket.AF_INET6, '::', 0, 0, 0)]
    with pytest.raises(Fatal) as excinfo:
        sshuttle.options.parse_subnetport('fc00::/129')
    assert str(excinfo.value) \
        == 'Slash in CIDR notation (/129) is not between 0 and 128'

def test_parse_subnetport_ip6_with_port():
    for ip_repr, ip in _ip6_reprs.items():
        assert sshuttle.options.parse_subnetport('[' + ip_repr + ']:80') \
            == [(socket.AF_INET6, ip, 128, 80, 80)]
        assert sshuttle.options.parse_subnetport('[' + ip_repr + ']:80-90') \
            == [(socket.AF_INET6, ip, 128, 80, 90)]

def test_parse_subnetport_ip6_with_mask_and_port():
    for ip_repr, ip in _ip6_reprs.items():
        assert sshuttle.options.parse_subnetport('[' + ip_repr + '/128]:80') \
            == [(socket.AF_INET6, ip, 128, 80, 80)]
        assert sshuttle.options.parse_subnetport('[' + ip_repr + '/16]:80-90')\
            == [(socket.AF_INET6, ip, 16, 80, 90)]

def test_convert_arg_line_to_args_skips_comments():
    parser = sshuttle.options.MyArgumentParser()
    assert parser.convert_arg_line_to_args("# whatever something") == []

@patch('sshuttle.options.socket.getaddrinfo', side_effect = _mock_getaddrinfo)
def test_parse_subnetport_host(mock_getaddrinfo):
    assert set(sshuttle.options.parse_subnetport('example.com')) \
        == set([
            (socket.AF_INET6, '2606:2800:220:1:248:1893:25c8:1946', 128, 0, 0),
            (socket.AF_INET, '93.184.216.34', 32, 0, 0),
        ])
    assert set(sshuttle.options.parse_subnetport('my.local')) \
        == set([
            (socket.AF_INET6, '::1', 128, 0, 0),
            (socket.AF_INET, '127.0.0.1', 32, 0, 0),
        ])

@patch('sshuttle.options.socket.getaddrinfo', side_effect = _mock_getaddrinfo)
def test_parse_subnetport_host_with_port(mock_getaddrinfo):
    assert set(sshuttle.options.parse_subnetport('example.com:80')) \
        == set([
            (socket.AF_INET6, '2606:2800:220:1:248:1893:25c8:1946', 128, 80, 80),
            (socket.AF_INET, '93.184.216.34', 32, 80, 80),
        ])
    assert set(sshuttle.options.parse_subnetport('example.com:80-90')) \
        == set([
            (socket.AF_INET6, '2606:2800:220:1:248:1893:25c8:1946', 128, 80, 90),
            (socket.AF_INET, '93.184.216.34', 32, 80, 90),
        ])
    assert set(sshuttle.options.parse_subnetport('my.local:445')) \
        == set([
            (socket.AF_INET6, '::1', 128, 445, 445),
            (socket.AF_INET, '127.0.0.1', 32, 445, 445),
        ])
    assert set(sshuttle.options.parse_subnetport('my.local:445-450')) \
        == set([
            (socket.AF_INET6, '::1', 128, 445, 450),
            (socket.AF_INET, '127.0.0.1', 32, 445, 450),
        ])
<|endoftext|>


<commit_before>"""
Tests for Advent of Code
"""
from . import day1, day2, day3, day4

def test_day_1():
    assert day1.solve('R2, L3') == 5
    assert day1.solve('R2, R2, R2') == 2
    assert day1.solve('R5, L5, R5, R3') == 12

    assert day1.solve('R8, R4, R4, R8', dupe=True) == 4

def test_day_2():
    assert day2.solve("ULL\nRRDDD\nLURDL\nUUUUD") == '1985'
    assert day2.solve("ULL\nRRDDD\nLURDL\nUUUUD", big=True) == '5DB3'

    assert day2.solve_better("ULL\nRRDDD\nLURDL\nUUUUD") == '1985'
    assert day2.solve_better("ULL\nRRDDD\nLURDL\nUUUUD", big=True) == '5DB3'

def test_day_3():
    assert day3.solve(["  5  10 25 "]) == 0
    assert day3.solve(["  25  20 10 "]) == 1

    assert day3.solve_vertical(['2 2 2', '2 3 4', '1 1 1', '10 11 12', '5 5 5', '6 6 6']) == 2

def test_day_4():
    assert day4.real_room('aaaaa-bbb-z-y-x', 'abxyz') is True
    assert day4.real_room('a-b-c-d-e-f-g-h', 'abcde') is True
    assert day4.real_room('not-a-real-room', 'oarel') is True
    assert day4.real_room('totally-real-room', 'decoy') is False

    assert day4.solve([
        'aaaaa-bbb-z-y-x-123[abxyz]',
        'a-b-c-d-e-f-g-h-987[abcde]',
        'not-a-real-room-404[oarel]',
        'totally-real-room-200[decoy]'
    ]) == 1514

    assert day4.rotate_room('qzmt-zixmtkozy-ivhz', 343) == 'very encrypted name'
<commit_msg>Tests: use pytest's parametrize<commit_after>"""
Tests for Advent of Code
"""
import pytest
from . import day1, day2, day3, day4

@pytest.mark.parametrize('directions,distance,dupe', [
    ('R2, L3', 5, False),
    ('R2, R2, R2', 2, False),
    ('R5, L5, R5, R3', 12, False),
    ('R8, R4, R4, R8', 4, True),
])
def test_day_1(directions, distance, dupe):
    assert day1.solve(directions, dupe=dupe) == distance

@pytest.mark.parametrize('directions,code,big', [
    ("ULL\nRRDDD\nLURDL\nUUUUD", '1985', False),
    ("ULL\nRRDDD\nLURDL\nUUUUD", '5DB3', True),
])
def test_day_2(directions, code, big):
    assert day2.solve(directions, big=big) == code
    assert day2.solve_better(directions, big=big) == code

@pytest.mark.parametrize('edges,count', [
    (["  5  10 25 "], 0),
    (["  25  20 10 "], 1),
])
def test_day_3(edges, count):
    assert day3.solve(edges) == count

def test_day_3_vertical():
    assert day3.solve_vertical(['2 2 2', '2 3 4', '1 1 1', '10 11 12', '5 5 5', '6 6 6']) == 2

@pytest.mark.parametrize('room,checksum,answer', [
    ('aaaaa-bbb-z-y-x', 'abxyz', True),
    ('a-b-c-d-e-f-g-h', 'abcde', True),
    ('not-a-real-room', 'oarel', True),
    ('totally-real-room', 'decoy', False),
])
def test_day_4_real_room(room, checksum, answer):
    assert day4.real_room(room, checksum) == answer

def test_day_4():
    assert day4.solve([
        'aaaaa-bbb-z-y-x-123[abxyz]',
        'a-b-c-d-e-f-g-h-987[abcde]',
        'not-a-real-room-404[oarel]',
        'totally-real-room-200[decoy]'
    ]) == 1514

    assert day4.rotate_room('qzmt-zixmtkozy-ivhz', 343) == 'very encrypted name'
<|endoftext|>


<commit_before># coding: utf8
from __future__ import unicode_literals
import os
import io
import sys

from nose.tools import *
from mock import patch, Mock
import newick
from clldutils.path import Path

from beastling.configuration import (
    Configuration, get_glottolog_data, _BEAST_MAX_LENGTH,
)
from beastling.beastxml import BeastXml
from .util import WithConfigAndTempDir, config_path

def check_lat_lon(provided, target_lat, target_lon):
    prov_lat, prov_lon = provided
    return round(prov_lat, 2) == target_lat and round(prov_lon, 2) == target_lon

class Tests(WithConfigAndTempDir):
    def _make_cfg(self, *names):
        return self.make_cfg([config_path(name).as_posix() for name in names])

    def _make_bad_cfg(self, name):
        return self.make_cfg(config_path(name, bad=True).as_posix())

    def test_get_glottolog_geo(self):
        geodata = self.tmp.joinpath('glottolog-2.5-geo.csv')
        with geodata.open('w', encoding='utf8') as fp:
            fp.write('x')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            self.assertEqual(Path(get_glottolog_data('geo', '2.5')), geodata)

    def test_get_glottolog_newick(self):
        with self.tmp.joinpath('glottolog-2.5.newick').open('w', encoding='utf8') as fp:
            fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            trees = newick.read(get_glottolog_data('newick', '2.5'))
            self.assertEqual(trees[0].name, 'A [abcd1234]')

    def test_get_glottolog_data_download(self):
        data_dir = os.path.join(self.tmp.as_posix(), 'data')

        class URLopener(object):
            def retrieve(self, url, fname):
                with io.open(fname, 'w', encoding='utf8') as fp:
                    fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        class URLopenerError(object):
            def retrieve(self, url, fname):
                raise IOError()

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopenerError,
        ):
            with self.assertRaises(ValueError):
                get_glottolog_data('newick', '2.5')

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopener,
        ):
            assert get_glottolog_data('newick', '2.5')

    def test_families(self):
        cfg1 = self._make_cfg('glottolog_families')
        cfg2 = self._make_cfg('glottolog_families_from_file')
        cfg1.process()
        cfg2.process()
        self.assertEqual(cfg1.languages, cfg2.languages)

    def test_config(self):
        cfg = Configuration(configfile={
            'admin': {

            },
            'languages': {
                'monophyly': True,
                'starting_tree': 'T',
                'sample_topology': False,
                'sample_branch_lengths': False,
            },
            'calibration': {
                'abcd1234, efgh5678': '10-20',
            },
            'model': {
                'binarised': True,
                'minimum_data': '4.5'
            },
            'model2': {
                'binarized': True,
            },
        })
        self.assertAlmostEqual(cfg.calibration_configs['abcd1234, efgh5678'], "10-20")
        self.assertTrue(cfg.model_configs[1]['binarised'])

        with self.assertRaisesRegexp(ValueError, 'Value for overlap') as e:
            Configuration(configfile={'languages': {'overlap': 'invalid'}, 'models': {}})

        with self.assertRaisesRegexp(ValueError, 'Config file') as e:
            Configuration(configfile={'languages': {}})

    @raises(ValueError)
    def test_no_data(self):
        cfg = self._make_bad_cfg("no_data")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_langs")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model_sec")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("unknown_model")
        cfg.process()

    @raises(ValueError)
    def test_bad_overlap(self):
        cfg = self._make_bad_cfg("bad_overlap")
        cfg.process()

    @raises(ValueError)
    def test_bad_frequencies(self):
        cfg = self._make_bad_cfg("bad_frequencies")
        cfg.process()
        # This is an error in a model section, which is only raised
        # when the model is constructed as XML.
        BeastXml(cfg)

    def test_calibration(self):
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertIn('Austronesian', config.calibrations)
        v = config.calibrations['Austronesian']
        xml1 = BeastXml(config).tostring().decode('utf8')
        print(config.calibrations)

        # Now remove one calibration point ...
        del config.calibrations['Austronesian']
        print(config.calibrations)
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertNotEqual(
            len(xml1.split('DistributionForAustronesianMRCA')),
            len(xml2.split('DistributionForAustronesianMRCA')))

        # ... and add it back in with using the glottocode:
        config.calibrations['aust1307'] = v
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertEqual(
            len(xml1.split('DistributionForAustronesianMRCA')),
            len(xml2.split('DistributionForaust1307MRCA')))

    def test_calibration_string_formats(self):
        # Test lower bound format
        config = self._make_cfg('basic', 'calibration_lower_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param[1], sys.maxsize)
        # Test upper bound format
        config = self._make_cfg('basic', 'calibration_upper_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param[0], 0.0)

        # Test range and param formats for all three distributions
        for dist in ('normal', 'lognormal', 'uniform'):
            for style in ('range', 'params'):
                config = self._make_cfg('basic', 'calibration_%s_%s' % (dist, style))
                config.process()
            self.assertEqual(list(config.calibrations.values())[0].dist, dist)

    @raises(ValueError)
    def test_calibration_bad_bounds(self):
        config = Configuration(configfile=[
            config_path("basic").as_posix(),
            config_path("bad_cal_endpoints", bad=True).as_posix(),
            ])
        config.process()

    @raises(ValueError)
    def test_calibration_bad_monophyly(self):
        config = Configuration(configfile=[
            config_path("basic").as_posix(),
            config_path("monophyletic").as_posix(),
            config_path("bad_cal_monophyly", bad=True).as_posix(),
            ])
        config.process()

    def test_overlong_chain(self):
        config = self._make_cfg('basic')
        config.chainlength = 9e999
        config.process()
        self.assertEqual(config.chainlength, _BEAST_MAX_LENGTH)

    def test_file_embedding(self):
        config = self._make_cfg('glottolog_families_from_file','embed_data')
        xml = BeastXml(config).tostring().decode('utf8')
        # Check for evidence of data
        self.assertTrue("aari1239,1,1,1,1,1,1,?,1,?,1" in xml)
        # Check for evidence of families
        self.assertTrue("Malayo-Polynesian" in xml)

    def test_minimum_data(self):
        # f8 has 60% missing data.  By default it should be included...
        config = self._make_cfg('basic')
        config.process()
        self.assertTrue("f8" in config.models[0].features)
        # ...but if we insist on 75% data or more it should be excluded...
        config = self._make_cfg('basic', 'minimum_data')
        config.process()
        self.assertEqual(config.models[0].minimum_data, 75.0)
        self.assertTrue("f8" not in config.models[0].features)

    def test_pruned_rlc(self):
        # Make sure pruned trees are disabled if used in conjunction with RLC
        config = self._make_cfg('basic', 'pruned', 'random')
        self.assertTrue(config.model_configs[0]["pruned"])
        config.process()
        self.assertFalse(config.models[0].pruned)

    def test_no_monophyly_geo(self):
        # Make sure that geographic sampling without monophyly constraints emits a warning
        config = self._make_cfg('basic', 'geo', 'geo_sampled')
        config.process()
        self.assertTrue(any(["[WARNING] Geographic sampling" in m for m in config.messages]))

    def test_ascertainment_auto_setting(self):
        # Without calibration, there should be no ascertainment...
        config = self._make_cfg('basic')
        config.process()
        self.assertFalse(config.models[0].ascertained)
        # But with it there should...
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # Unless, of course, we have constant data...
        config = self._make_cfg('covarion_multistate', 'calibration')
        config.model_configs[0]["remove_constant_features"] = False
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_ascertainment_override(self):
        # Make sure we can override the above automagic
        config = self._make_cfg('basic', 'ascertainment_true')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # And with calibrations...
        config = self._make_cfg('basic', 'calibration', 'ascertainment_false')
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_bad_ascertainment(self):
        # Make sure we refuse to produce a misspecified model
        config = self._make_cfg('covarion_multistate','ascertainment_true')
        config.model_configs[0]["remove_constant_features"] = False
        with self.assertRaises(ValueError):
            config.process()

    def test_binarisation_ascertainment(self):
        # Even with ascertainment = False, ascertainment should still
        # be done for recoded data
        config = self._make_cfg('covarion_binarised','ascertainment_false')
        config.process()
        xml = BeastXml(config).tostring().decode('utf8')
        self.assertTrue("ascertained=\"true\"" in xml)
        self.assertTrue("excludeto=\"1\"" in xml)

    def test_user_locations(self):
        # First check that we correctly load Glottolog's locations for aiw and abp
        config = self._make_cfg('basic', 'geo')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 5.95, 36.57))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Now check that we can overwrite just one of these...
        config = self._make_cfg('basic', 'geo', 'geo_user_loc')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Make sure that specifying the location data in [languages] caused a deprecation warning
        self.assertTrue(len(config.urgent_messages) > 0)
        # Repeat the above test but specifying location data in [geography], which should cause no warning
        config = self._make_cfg('basic', 'geo', 'new_geo_user_loc')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        self.assertTrue(len(config.urgent_messages) == 0)
        # Now check that we can overwrite them both using multiple files
        config = self._make_cfg('basic', 'geo', 'geo_user_loc_multifile')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 6.66, 6.66))

    def test_monophyly_levels(self):
        # The isolates.csv data file contains Japanese, Korean and Basque, plus
        # English and Russian.  When used with standard monophly, we should see
        # a four-way polytomy with eng+rus grouped (IE) and the rest isolated.
        config = self._make_cfg('admin', 'mk', 'isolates', 'monophyletic')
        config.process()
        tree = newick.loads(config.monophyly_newick)[0]
        assert len(tree.descendants) == 4
        for node in tree.descendants:
            if len(node.descendants) == 2:
                assert all((l.is_leaf and l.name in ("eng", "rus") for l in node.descendants))
        # Now we set monophyly_start_depth to 1, i.e. ignore the top-most
        # level of Glottolog constraints.  Now we should just have a massive
        # polytomy, since IE no longer matters and eng and rus are in separate
        # subfamilies (Germanic vs Balto-Slavic).
        config = self._make_cfg('admin', 'mk', 'isolates', 'monophyletic-start-depth')
        config.process()
        tree = newick.loads(config.monophyly_newick)[0]
        assert len(tree.descendants) == 5

    def test_subsampling(self):
        # First check how many languages there usually are
        config = self._make_cfg('admin', 'mk')
        config.process()
        full_lang_count = len(config.languages)
        # Try various subsamples and make sure they work
        for subsample_size in range(2, full_lang_count):
            config = self._make_cfg('admin', 'mk')
            config.subsample_size = subsample_size
            config.process()
            assert len(config.languages) == subsample_size
        # Make sure if we ask for more languages than we have nothing happens
        config = self._make_cfg('admin', 'mk')
        config.subsample_size = full_lang_count + 42
        config.process()
        assert len(config.languages) == full_lang_count

    def test_language_groups(self):
        config = self._make_cfg('basic', 'taxa')
        config.process()
        self.assertEqual(config.language_groups["abf"], {"abf"})
        self.assertEqual(config.language_groups["macronesian"], {"kbt", "abf", "abg"})

    @raises(KeyError)
    def test_nonexisting_language_group(self):
        config = self._make_cfg('basic', 'reconstruct_one')
        config.process()
<commit_msg>Remove debug output<commit_after># coding: utf8
from __future__ import unicode_literals
import os
import io
import sys

from nose.tools import *
from mock import patch, Mock
import newick
from clldutils.path import Path

from beastling.configuration import (
    Configuration, get_glottolog_data, _BEAST_MAX_LENGTH,
)
from beastling.beastxml import BeastXml
from .util import WithConfigAndTempDir, config_path

def check_lat_lon(provided, target_lat, target_lon):
    prov_lat, prov_lon = provided
    return round(prov_lat, 2) == target_lat and round(prov_lon, 2) == target_lon

class Tests(WithConfigAndTempDir):
    def _make_cfg(self, *names):
        return self.make_cfg([config_path(name).as_posix() for name in names])

    def _make_bad_cfg(self, name):
        return self.make_cfg(config_path(name, bad=True).as_posix())

    def test_get_glottolog_geo(self):
        geodata = self.tmp.joinpath('glottolog-2.5-geo.csv')
        with geodata.open('w', encoding='utf8') as fp:
            fp.write('x')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            self.assertEqual(Path(get_glottolog_data('geo', '2.5')), geodata)

    def test_get_glottolog_newick(self):
        with self.tmp.joinpath('glottolog-2.5.newick').open('w', encoding='utf8') as fp:
            fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            trees = newick.read(get_glottolog_data('newick', '2.5'))
            self.assertEqual(trees[0].name, 'A [abcd1234]')

    def test_get_glottolog_data_download(self):
        data_dir = os.path.join(self.tmp.as_posix(), 'data')

        class URLopener(object):
            def retrieve(self, url, fname):
                with io.open(fname, 'w', encoding='utf8') as fp:
                    fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        class URLopenerError(object):
            def retrieve(self, url, fname):
                raise IOError()

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopenerError,
        ):
            with self.assertRaises(ValueError):
                get_glottolog_data('newick', '2.5')

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopener,
        ):
            assert get_glottolog_data('newick', '2.5')

    def test_families(self):
        cfg1 = self._make_cfg('glottolog_families')
        cfg2 = self._make_cfg('glottolog_families_from_file')
        cfg1.process()
        cfg2.process()
        self.assertEqual(cfg1.languages, cfg2.languages)

    def test_config(self):
        cfg = Configuration(configfile={
            'admin': {

            },
            'languages': {
                'monophyly': True,
                'starting_tree': 'T',
                'sample_topology': False,
                'sample_branch_lengths': False,
            },
            'calibration': {
                'abcd1234, efgh5678': '10-20',
            },
            'model': {
                'binarised': True,
                'minimum_data': '4.5'
            },
            'model2': {
                'binarized': True,
            },
        })
        self.assertAlmostEqual(cfg.calibration_configs['abcd1234, efgh5678'], "10-20")
        self.assertTrue(cfg.model_configs[1]['binarised'])

        with self.assertRaisesRegexp(ValueError, 'Value for overlap') as e:
            Configuration(configfile={'languages': {'overlap': 'invalid'}, 'models': {}})

        with self.assertRaisesRegexp(ValueError, 'Config file') as e:
            Configuration(configfile={'languages': {}})

    @raises(ValueError)
    def test_no_data(self):
        cfg = self._make_bad_cfg("no_data")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_langs")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model_sec")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("unknown_model")
        cfg.process()

    @raises(ValueError)
    def test_bad_overlap(self):
        cfg = self._make_bad_cfg("bad_overlap")
        cfg.process()

    @raises(ValueError)
    def test_bad_frequencies(self):
        cfg = self._make_bad_cfg("bad_frequencies")
        cfg.process()
        # This is an error in a model section, which is only raised
        # when the model is constructed as XML.
        BeastXml(cfg)

    def test_calibration(self):
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertIn('Austronesian', config.calibrations)
        v = config.calibrations['Austronesian']
        xml1 = BeastXml(config).tostring().decode('utf8')

        # Now remove one calibration point ...
        del config.calibrations['Austronesian']
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertNotEqual(
            len(xml1.split('DistributionForAustronesianMRCA')),
            len(xml2.split('DistributionForAustronesianMRCA')))

        # ... and add it back in with using the glottocode:
        config.calibrations['aust1307'] = v
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertEqual(
            len(xml1.split('DistributionForAustronesianMRCA')),
            len(xml2.split('DistributionForaust1307MRCA')))

    def test_calibration_string_formats(self):
        # Test lower bound format
        config = self._make_cfg('basic', 'calibration_lower_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param[1], sys.maxsize)
        # Test upper bound format
        config = self._make_cfg('basic', 'calibration_upper_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param[0], 0.0)

        # Test range and param formats for all three distributions
        for dist in ('normal', 'lognormal', 'uniform'):
            for style in ('range', 'params'):
                config = self._make_cfg('basic', 'calibration_%s_%s' % (dist, style))
                config.process()
            self.assertEqual(list(config.calibrations.values())[0].dist, dist)

    @raises(ValueError)
    def test_calibration_bad_bounds(self):
        config = Configuration(configfile=[
            config_path("basic").as_posix(),
            config_path("bad_cal_endpoints", bad=True).as_posix(),
            ])
        config.process()

    @raises(ValueError)
    def test_calibration_bad_monophyly(self):
        config = Configuration(configfile=[
            config_path("basic").as_posix(),
            config_path("monophyletic").as_posix(),
            config_path("bad_cal_monophyly", bad=True).as_posix(),
            ])
        config.process()

    def test_overlong_chain(self):
        config = self._make_cfg('basic')
        config.chainlength = 9e999
        config.process()
        self.assertEqual(config.chainlength, _BEAST_MAX_LENGTH)

    def test_file_embedding(self):
        config = self._make_cfg('glottolog_families_from_file','embed_data')
        xml = BeastXml(config).tostring().decode('utf8')
        # Check for evidence of data
        self.assertTrue("aari1239,1,1,1,1,1,1,?,1,?,1" in xml)
        # Check for evidence of families
        self.assertTrue("Malayo-Polynesian" in xml)

    def test_minimum_data(self):
        # f8 has 60% missing data.  By default it should be included...
        config = self._make_cfg('basic')
        config.process()
        self.assertTrue("f8" in config.models[0].features)
        # ...but if we insist on 75% data or more it should be excluded...
        config = self._make_cfg('basic', 'minimum_data')
        config.process()
        self.assertEqual(config.models[0].minimum_data, 75.0)
        self.assertTrue("f8" not in config.models[0].features)

    def test_pruned_rlc(self):
        # Make sure pruned trees are disabled if used in conjunction with RLC
        config = self._make_cfg('basic', 'pruned', 'random')
        self.assertTrue(config.model_configs[0]["pruned"])
        config.process()
        self.assertFalse(config.models[0].pruned)

    def test_no_monophyly_geo(self):
        # Make sure that geographic sampling without monophyly constraints emits a warning
        config = self._make_cfg('basic', 'geo', 'geo_sampled')
        config.process()
        self.assertTrue(any(["[WARNING] Geographic sampling" in m for m in config.messages]))

    def test_ascertainment_auto_setting(self):
        # Without calibration, there should be no ascertainment...
        config = self._make_cfg('basic')
        config.process()
        self.assertFalse(config.models[0].ascertained)
        # But with it there should...
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # Unless, of course, we have constant data...
        config = self._make_cfg('covarion_multistate', 'calibration')
        config.model_configs[0]["remove_constant_features"] = False
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_ascertainment_override(self):
        # Make sure we can override the above automagic
        config = self._make_cfg('basic', 'ascertainment_true')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # And with calibrations...
        config = self._make_cfg('basic', 'calibration', 'ascertainment_false')
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_bad_ascertainment(self):
        # Make sure we refuse to produce a misspecified model
        config = self._make_cfg('covarion_multistate','ascertainment_true')
        config.model_configs[0]["remove_constant_features"] = False
        with self.assertRaises(ValueError):
            config.process()

    def test_binarisation_ascertainment(self):
        # Even with ascertainment = False, ascertainment should still
        # be done for recoded data
        config = self._make_cfg('covarion_binarised','ascertainment_false')
        config.process()
        xml = BeastXml(config).tostring().decode('utf8')
        self.assertTrue("ascertained=\"true\"" in xml)
        self.assertTrue("excludeto=\"1\"" in xml)

    def test_user_locations(self):
        # First check that we correctly load Glottolog's locations for aiw and abp
        config = self._make_cfg('basic', 'geo')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 5.95, 36.57))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Now check that we can overwrite just one of these...
        config = self._make_cfg('basic', 'geo', 'geo_user_loc')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Make sure that specifying the location data in [languages] caused a deprecation warning
        self.assertTrue(len(config.urgent_messages) > 0)
        # Repeat the above test but specifying location data in [geography], which should cause no warning
        config = self._make_cfg('basic', 'geo', 'new_geo_user_loc')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        self.assertTrue(len(config.urgent_messages) == 0)
        # Now check that we can overwrite them both using multiple files
        config = self._make_cfg('basic', 'geo', 'geo_user_loc_multifile')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 6.66, 6.66))

    def test_monophyly_levels(self):
        # The isolates.csv data file contains Japanese, Korean and Basque, plus
        # English and Russian.  When used with standard monophly, we should see
        # a four-way polytomy with eng+rus grouped (IE) and the rest isolated.
        config = self._make_cfg('admin', 'mk', 'isolates', 'monophyletic')
        config.process()
        tree = newick.loads(config.monophyly_newick)[0]
        assert len(tree.descendants) == 4
        for node in tree.descendants:
            if len(node.descendants) == 2:
                assert all((l.is_leaf and l.name in ("eng", "rus") for l in node.descendants))
        # Now we set monophyly_start_depth to 1, i.e. ignore the top-most
        # level of Glottolog constraints.  Now we should just have a massive
        # polytomy, since IE no longer matters and eng and rus are in separate
        # subfamilies (Germanic vs Balto-Slavic).
        config = self._make_cfg('admin', 'mk', 'isolates', 'monophyletic-start-depth')
        config.process()
        tree = newick.loads(config.monophyly_newick)[0]
        assert len(tree.descendants) == 5

    def test_subsampling(self):
        # First check how many languages there usually are
        config = self._make_cfg('admin', 'mk')
        config.process()
        full_lang_count = len(config.languages)
        # Try various subsamples and make sure they work
        for subsample_size in range(2, full_lang_count):
            config = self._make_cfg('admin', 'mk')
            config.subsample_size = subsample_size
            config.process()
            assert len(config.languages) == subsample_size
        # Make sure if we ask for more languages than we have nothing happens
        config = self._make_cfg('admin', 'mk')
        config.subsample_size = full_lang_count + 42
        config.process()
        assert len(config.languages) == full_lang_count

    def test_language_groups(self):
        config = self._make_cfg('basic', 'taxa')
        config.process()
        self.assertEqual(config.language_groups["abf"], {"abf"})
        self.assertEqual(config.language_groups["macronesian"], {"kbt", "abf", "abg"})

    @raises(KeyError)
    def test_nonexisting_language_group(self):
        config = self._make_cfg('basic', 'reconstruct_one')
        config.process()
<|endoftext|>


<commit_before>import networkx

import pyvex
import simuvex

class Blade(object):
    """
    Blade is a light-weight program slicer that works with networkx DiGraph containing SimIRSBs.
    It is meant to be used in angr for small or on-the-fly analyses.
    """
    def __init__(self, graph, dst_run, dst_stmt_idx, direction='backward', project=None, cfg=None, ignore_sp=False,
                 ignore_bp=False, ignored_regs=None, max_level=3):
        """
        :param networkx.DiGraph graph:  A graph representing the control flow graph. Note that it does not take
                                        angr.analyses.CFGAccurate or angr.analyses.CFGFast.
        :param int dst_run:             An address specifying the target SimRun.
        :param int dst_stmt_idx:        The target statement index. -1 means executing until the last statement.
        :param str direction:           'backward' or 'forward' slicing. Forward slicing is not yet supported.
        :param angr.Project project:    The project instance.
        :param angr.analyses.CFGBase cfg: the CFG instance. It will be made mandatory later.
        :param bool ignore_sp:          Whether the stack pointer should be ignored in dependency tracking. Any
                                        dependency from/to stack pointers will be ignored if this options is True.
        :param ignore_bp:               Whether the base pointer should be ignored or not.
        :return: None
        """

        self._graph = graph
        self._dst_run = dst_run
        self._dst_stmt_idx = dst_stmt_idx
        self._ignore_sp = ignore_sp
        self._ignore_bp = ignore_bp
        self._max_level = max_level

        self._slice = networkx.DiGraph()

        self.project = project
        self._cfg = cfg
        if self._cfg is None:
            # `cfg` is made optional only for compatibility concern. It will be made a positional parameter later.
            raise AngrBladeError('"cfg" must be specified.')

        if not self._in_graph(self._dst_run):
            raise AngrBladeError("The specified SimRun %s doesn't exist in graph.")

        self._ignored_regs = set()
        if ignored_regs:
            for r in ignored_regs:
                if isinstance(r, (int, long)):
                    self._ignored_regs.add(r)
                else:
                    self._ignored_regs.add(self.project.arch.registers[r][0])

        self._run_cache = { }

        self._traced_runs = set()

        if direction == 'backward':
            self._backward_slice()
        elif direction == 'forward':
            raise AngrBladeError('Forward slicing is not implemented yet')
        else:
            raise AngrBladeError("Unknown slicing direction %s", direction)

    #
    # Properties
    #

    @property
    def slice(self):
        return self._slice

    #
    # Public methods
    #

    def dbg_repr(self):
        s = ""

        block_addrs = list(set([ a for a, _ in self.slice.nodes_iter() ]))

        for block_addr in block_addrs:
            block_str = "IRSB %08x\n" % block_addr

            block = self.project.factory.block(block_addr).vex

            included_stmts = set([ stmt for _, stmt in self.slice.nodes_iter() if _ == block_addr ])

            for i, stmt in enumerate(block.statements):
                block_str += "%02s: %s\n" % ("+" if i in included_stmts else "-",
                                   str(stmt)
                                   )

            s += block_str
            s += "\n"

        return s

    #
    # Private methods
    #

    def _get_irsb(self, v):
        """
        Get the IRSB object from an address, a simuvex.SimProcedure, a simuvex.SimIRSB, or a CFGNode.
        :param v: Can be one of the following: an address, a simuvex.SimProcedure, a simuvex.SimIRSB, or a CFGNode.
        :return: The IRSB instance.
        :rtype: pyvex.IRSB
        """

        if isinstance(v, simuvex.SimProcedure):
            raise AngrBladeSimProcError()

        elif isinstance(v, simuvex.SimIRSB):
            v = v.addr

        elif isinstance(v, CFGNode):
            v = v.addr

        if type(v) in (int, long):
            # Generate an IRSB from self._project

            if v in self._run_cache:
                return self._run_cache[v]

            if self.project:
                irsb = self.project.factory.block(v).vex
                self._run_cache[v] = irsb
                return irsb
            else:
                raise AngrBladeError("Project must be specified if you give me all addresses for SimRuns")

        else:
            raise AngrBladeError('Unsupported SimRun argument type %s', type(v))

    def _get_cfgnode(self, thing):
        """
        Get the CFGNode corresponding to the specific address.

        :param thing: Can be anything that self._normalize() accepts. Usually it's the address of the node
        :return: the CFGNode instance
        :rtype: CFGNode
        """

        return self._cfg.get_any_node(self._get_addr(thing))

    def _get_addr(self, v):
        """
        Get address of the basic block or CFG node specified by v.
        :param v: Can be one of the following: a simuvex.SimIRSB, a simuvex.SimProcedure, a CFGNode, or an address.
        :return: The address.
        :rtype: int
        """

        if isinstance(v, simuvex.SimIRSB) or isinstance(v, simuvex.SimProcedure):
            if type(self._graph.nodes()[0]) in (int, long):
                return v.addr
            else:
                return v
        elif isinstance(v, CFGNode):
            return v.addr
        elif type(v) in (int, long):
            return v
        else:
            raise AngrBladeError('Unsupported SimRun argument type %s', type(v))

    def _in_graph(self, v):
        return self._get_cfgnode(v) in self._graph

    def _inslice_callback(self, stmt_idx, stmt, infodict):
        tpl = (infodict['irsb_addr'], stmt_idx)
        if 'prev' in infodict and infodict['prev']:
            prev = infodict['prev']
            self._slice.add_edge(tpl, prev)
        else:
            self._slice.add_node(tpl)

        infodict['prev'] = tpl

    def _backward_slice(self):
        """
        Backward slicing.

        We support the following IRStmts:
        # WrTmp
        # Put

        We support the following IRExprs:
        # Get
        # RdTmp
        # Const

        :return:
        """

        temps = set()
        regs = set()

        # Retrieve the target: are we slicing from a register(IRStmt.Put), or a temp(IRStmt.WrTmp)?
        stmts = self._get_irsb(self._dst_run).statements

        if self._dst_stmt_idx != -1:
            dst_stmt = stmts[self._dst_stmt_idx]

            if type(dst_stmt) is pyvex.IRStmt.Put:
                regs.add(dst_stmt.offset)
            elif type(dst_stmt) is pyvex.IRStmt.WrTmp:
                temps.add(dst_stmt.tmp)
            else:
                raise AngrBladeError('Incorrect type of the specified target statement. We only support Put and WrTmp.')

            prev = (self._get_addr(self._dst_run), self._dst_stmt_idx)
        else:

            next_expr = self._get_irsb(self._dst_run).next

            if type(next_expr) is pyvex.IRExpr.RdTmp:
                temps.add(next_expr.tmp)
            elif type(next_expr) is pyvex.IRExpr.Const:
                # A const doesn't rely on anything else!
                pass
            else:
                raise AngrBladeError('Unsupported type for irsb.next: %s' % type(next_expr))

            # Then we gotta start from the very last statement!
            self._dst_stmt_idx = len(stmts) - 1

            prev = (self._get_addr(self._dst_run), 'default')

        slicer = simuvex.SimSlicer(self.project.arch, stmts,
                                   target_tmps=temps,
                                   target_regs=regs,
                                   target_stack_offsets=None,
                                   inslice_callback=self._inslice_callback,
                                   inslice_callback_infodict={
                                       'irsb_addr':  self._get_irsb(self._dst_run)._addr,
                                       'prev': prev,
                                   })
        regs = slicer.final_regs
        if self._ignore_sp and self.project.arch.sp_offset in regs:
            regs.remove(self.project.arch.sp_offset)
        if self._ignore_bp and self.project.arch.bp_offset in regs:
            regs.remove(self.project.arch.bp_offset)
        for offset in self._ignored_regs:
            if offset in regs:
                regs.remove(offset)

        stack_offsets = slicer.final_stack_offsets

        prev = slicer.inslice_callback_infodict['prev']

        if regs or stack_offsets:
            cfgnode = self._get_cfgnode(self._dst_run)
            in_edges = self._graph.in_edges(cfgnode, data=True)

            for pred, _, data in in_edges:
                if pred not in self._traced_runs:
                    self._traced_runs.add(pred)
                    self._backward_slice_recursive(self._max_level - 1, pred, regs, stack_offsets, prev, data['stmt_idx'])

    def _backward_slice_recursive(self, level, run, regs, stack_offsets, prev, exit_stmt_idx):

        if level <= 0:
            return

        temps = set()
        regs = regs.copy()

        stmts = self._get_irsb(run).statements

        if exit_stmt_idx is None or exit_stmt_idx == 'default':
            # Initialize the temps set with whatever in the `next` attribute of this irsb
            next_expr = self._get_irsb(run).next
            if type(next_expr) is pyvex.IRExpr.RdTmp:
                temps.add(next_expr.tmp)

        else:
            exit_stmt = self._get_irsb(run).statements[exit_stmt_idx]

            if type(exit_stmt.guard) is pyvex.IRExpr.RdTmp:
                temps.add(exit_stmt.guard.tmp)

            # Put it in our slice
            irsb_addr = self._get_addr(run)
            self._inslice_callback(exit_stmt_idx, exit_stmt, {'irsb_addr': irsb_addr, 'prev': prev})
            prev = (irsb_addr, exit_stmt_idx)

        slicer = simuvex.SimSlicer(self.project.arch, stmts,
                                   target_tmps=temps,
                                   target_regs=regs,
                                   target_stack_offsets=stack_offsets,
                                   inslice_callback=self._inslice_callback,
                                   inslice_callback_infodict={
                                       'irsb_addr' : self._get_addr(run),
                                       'prev' : prev
                                   })
        regs = slicer.final_regs

        if self._ignore_sp and self.project.arch.sp_offset in regs:
            regs.remove(self.project.arch.sp_offset)
        if self._ignore_bp and self.project.arch.bp_offset in regs:
            regs.remove(self.project.arch.bp_offset)

        stack_offsets = slicer.final_stack_offsets

        prev = slicer.inslice_callback_infodict['prev']

        if regs or stack_offsets:
            in_edges = self._graph.in_edges(self._get_cfgnode(run), data=True)

            for pred, _, data in in_edges:
                if pred not in self._traced_runs:
                    self._traced_runs.add(pred)
                    self._backward_slice_recursive(level - 1, pred, regs, stack_offsets, prev, data['stmt_idx'])

from .errors import AngrBladeError, AngrBladeSimProcError
from .analyses.cfg_node import CFGNode
<commit_msg>more comments for Blade.<commit_after>import networkx

import pyvex
import simuvex

class Blade(object):
    """
    Blade is a light-weight program slicer that works with networkx DiGraph containing SimIRSBs.
    It is meant to be used in angr for small or on-the-fly analyses.
    """
    def __init__(self, graph, dst_run, dst_stmt_idx, direction='backward', project=None, cfg=None, ignore_sp=False,
                 ignore_bp=False, ignored_regs=None, max_level=3):
        """
        :param networkx.DiGraph graph:  A graph representing the control flow graph. Note that it does not take
                                        angr.analyses.CFGAccurate or angr.analyses.CFGFast.
        :param int dst_run:             An address specifying the target SimRun.
        :param int dst_stmt_idx:        The target statement index. -1 means executing until the last statement.
        :param str direction:           'backward' or 'forward' slicing. Forward slicing is not yet supported.
        :param angr.Project project:    The project instance.
        :param angr.analyses.CFGBase cfg: the CFG instance. It will be made mandatory later.
        :param bool ignore_sp:          Whether the stack pointer should be ignored in dependency tracking. Any
                                        dependency from/to stack pointers will be ignored if this options is True.
        :param bool ignore_bp:          Whether the base pointer should be ignored or not.
        :param int  max_level:          The maximum number of blocks that we trace back for.
        :return: None
        """

        self._graph = graph
        self._dst_run = dst_run
        self._dst_stmt_idx = dst_stmt_idx
        self._ignore_sp = ignore_sp
        self._ignore_bp = ignore_bp
        self._max_level = max_level

        self._slice = networkx.DiGraph()

        self.project = project
        self._cfg = cfg
        if self._cfg is None:
            # `cfg` is made optional only for compatibility concern. It will be made a positional parameter later.
            raise AngrBladeError('"cfg" must be specified.')

        if not self._in_graph(self._dst_run):
            raise AngrBladeError("The specified SimRun %s doesn't exist in graph.")

        self._ignored_regs = set()
        if ignored_regs:
            for r in ignored_regs:
                if isinstance(r, (int, long)):
                    self._ignored_regs.add(r)
                else:
                    self._ignored_regs.add(self.project.arch.registers[r][0])

        self._run_cache = { }

        self._traced_runs = set()

        if direction == 'backward':
            self._backward_slice()
        elif direction == 'forward':
            raise AngrBladeError('Forward slicing is not implemented yet')
        else:
            raise AngrBladeError("Unknown slicing direction %s", direction)

    #
    # Properties
    #

    @property
    def slice(self):
        return self._slice

    #
    # Public methods
    #

    def dbg_repr(self):
        s = ""

        block_addrs = list(set([ a for a, _ in self.slice.nodes_iter() ]))

        for block_addr in block_addrs:
            block_str = "IRSB %08x\n" % block_addr

            block = self.project.factory.block(block_addr).vex

            included_stmts = set([ stmt for _, stmt in self.slice.nodes_iter() if _ == block_addr ])

            for i, stmt in enumerate(block.statements):
                block_str += "%02s: %s\n" % ("+" if i in included_stmts else "-",
                                   str(stmt)
                                   )

            s += block_str
            s += "\n"

        return s

    #
    # Private methods
    #

    def _get_irsb(self, v):
        """
        Get the IRSB object from an address, a simuvex.SimProcedure, a simuvex.SimIRSB, or a CFGNode.
        :param v: Can be one of the following: an address, a simuvex.SimProcedure, a simuvex.SimIRSB, or a CFGNode.
        :return: The IRSB instance.
        :rtype: pyvex.IRSB
        """

        if isinstance(v, simuvex.SimProcedure):
            raise AngrBladeSimProcError()

        elif isinstance(v, simuvex.SimIRSB):
            v = v.addr

        elif isinstance(v, CFGNode):
            v = v.addr

        if type(v) in (int, long):
            # Generate an IRSB from self._project

            if v in self._run_cache:
                return self._run_cache[v]

            if self.project:
                irsb = self.project.factory.block(v).vex
                self._run_cache[v] = irsb
                return irsb
            else:
                raise AngrBladeError("Project must be specified if you give me all addresses for SimRuns")

        else:
            raise AngrBladeError('Unsupported SimRun argument type %s', type(v))

    def _get_cfgnode(self, thing):
        """
        Get the CFGNode corresponding to the specific address.

        :param thing: Can be anything that self._normalize() accepts. Usually it's the address of the node
        :return: the CFGNode instance
        :rtype: CFGNode
        """

        return self._cfg.get_any_node(self._get_addr(thing))

    def _get_addr(self, v):
        """
        Get address of the basic block or CFG node specified by v.
        :param v: Can be one of the following: a simuvex.SimIRSB, a simuvex.SimProcedure, a CFGNode, or an address.
        :return: The address.
        :rtype: int
        """

        if isinstance(v, simuvex.SimIRSB) or isinstance(v, simuvex.SimProcedure):
            if type(self._graph.nodes()[0]) in (int, long):
                return v.addr
            else:
                return v
        elif isinstance(v, CFGNode):
            return v.addr
        elif type(v) in (int, long):
            return v
        else:
            raise AngrBladeError('Unsupported SimRun argument type %s', type(v))

    def _in_graph(self, v):
        return self._get_cfgnode(v) in self._graph

    def _inslice_callback(self, stmt_idx, stmt, infodict):
        tpl = (infodict['irsb_addr'], stmt_idx)
        if 'prev' in infodict and infodict['prev']:
            prev = infodict['prev']
            self._slice.add_edge(tpl, prev)
        else:
            self._slice.add_node(tpl)

        infodict['prev'] = tpl

    def _backward_slice(self):
        """
        Backward slicing.

        We support the following IRStmts:
        # WrTmp
        # Put

        We support the following IRExprs:
        # Get
        # RdTmp
        # Const

        :return:
        """

        temps = set()
        regs = set()

        # Retrieve the target: are we slicing from a register(IRStmt.Put), or a temp(IRStmt.WrTmp)?
        stmts = self._get_irsb(self._dst_run).statements

        if self._dst_stmt_idx != -1:
            dst_stmt = stmts[self._dst_stmt_idx]

            if type(dst_stmt) is pyvex.IRStmt.Put:
                regs.add(dst_stmt.offset)
            elif type(dst_stmt) is pyvex.IRStmt.WrTmp:
                temps.add(dst_stmt.tmp)
            else:
                raise AngrBladeError('Incorrect type of the specified target statement. We only support Put and WrTmp.')

            prev = (self._get_addr(self._dst_run), self._dst_stmt_idx)
        else:

            next_expr = self._get_irsb(self._dst_run).next

            if type(next_expr) is pyvex.IRExpr.RdTmp:
                temps.add(next_expr.tmp)
            elif type(next_expr) is pyvex.IRExpr.Const:
                # A const doesn't rely on anything else!
                pass
            else:
                raise AngrBladeError('Unsupported type for irsb.next: %s' % type(next_expr))

            # Then we gotta start from the very last statement!
            self._dst_stmt_idx = len(stmts) - 1

            prev = (self._get_addr(self._dst_run), 'default')

        slicer = simuvex.SimSlicer(self.project.arch, stmts,
                                   target_tmps=temps,
                                   target_regs=regs,
                                   target_stack_offsets=None,
                                   inslice_callback=self._inslice_callback,
                                   inslice_callback_infodict={
                                       'irsb_addr':  self._get_irsb(self._dst_run)._addr,
                                       'prev': prev,
                                   })
        regs = slicer.final_regs
        if self._ignore_sp and self.project.arch.sp_offset in regs:
            regs.remove(self.project.arch.sp_offset)
        if self._ignore_bp and self.project.arch.bp_offset in regs:
            regs.remove(self.project.arch.bp_offset)
        for offset in self._ignored_regs:
            if offset in regs:
                regs.remove(offset)

        stack_offsets = slicer.final_stack_offsets

        prev = slicer.inslice_callback_infodict['prev']

        if regs or stack_offsets:
            cfgnode = self._get_cfgnode(self._dst_run)
            in_edges = self._graph.in_edges(cfgnode, data=True)

            for pred, _, data in in_edges:
                if pred not in self._traced_runs:
                    self._traced_runs.add(pred)
                    self._backward_slice_recursive(self._max_level - 1, pred, regs, stack_offsets, prev, data['stmt_idx'])

    def _backward_slice_recursive(self, level, run, regs, stack_offsets, prev, exit_stmt_idx):

        if level <= 0:
            return

        temps = set()
        regs = regs.copy()

        stmts = self._get_irsb(run).statements

        if exit_stmt_idx is None or exit_stmt_idx == 'default':
            # Initialize the temps set with whatever in the `next` attribute of this irsb
            next_expr = self._get_irsb(run).next
            if type(next_expr) is pyvex.IRExpr.RdTmp:
                temps.add(next_expr.tmp)

        else:
            exit_stmt = self._get_irsb(run).statements[exit_stmt_idx]

            if type(exit_stmt.guard) is pyvex.IRExpr.RdTmp:
                temps.add(exit_stmt.guard.tmp)

            # Put it in our slice
            irsb_addr = self._get_addr(run)
            self._inslice_callback(exit_stmt_idx, exit_stmt, {'irsb_addr': irsb_addr, 'prev': prev})
            prev = (irsb_addr, exit_stmt_idx)

        slicer = simuvex.SimSlicer(self.project.arch, stmts,
                                   target_tmps=temps,
                                   target_regs=regs,
                                   target_stack_offsets=stack_offsets,
                                   inslice_callback=self._inslice_callback,
                                   inslice_callback_infodict={
                                       'irsb_addr' : self._get_addr(run),
                                       'prev' : prev
                                   })
        regs = slicer.final_regs

        if self._ignore_sp and self.project.arch.sp_offset in regs:
            regs.remove(self.project.arch.sp_offset)
        if self._ignore_bp and self.project.arch.bp_offset in regs:
            regs.remove(self.project.arch.bp_offset)

        stack_offsets = slicer.final_stack_offsets

        prev = slicer.inslice_callback_infodict['prev']

        if regs or stack_offsets:
            in_edges = self._graph.in_edges(self._get_cfgnode(run), data=True)

            for pred, _, data in in_edges:
                if pred not in self._traced_runs:
                    self._traced_runs.add(pred)
                    self._backward_slice_recursive(level - 1, pred, regs, stack_offsets, prev, data['stmt_idx'])

from .errors import AngrBladeError, AngrBladeSimProcError
from .analyses.cfg_node import CFGNode
<|endoftext|>