<commit_before>import os
import sys
import unittest

import DistRDF
import pyspark
from DistRDF.Backends import Spark
from DistRDF.Backends.Spark import Backend

class SparkBackendInitTest(unittest.TestCase):
    """
    Tests to ensure that the instance variables
    of `Spark` class are set according to the
    input `config` dict.
    """

    @classmethod
    def setUpClass(cls):
        """
        Synchronize PYSPARK_PYTHON variable to the current Python executable.

        Needed to avoid mismatch between python versions on driver and on
        the fake executor on the same machine.
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    @classmethod
    def tearDownClass(cls):
        """
        Stop the SparkContext and reset environment variable.
        """
        os.environ["PYSPARK_PYTHON"] = ""

    def test_set_spark_context_default(self):
        """
        Check that a `SparkContext` object is created with default options for
        the current system.
        """
        backend = Backend.SparkBackend()

        self.assertIsInstance(backend.sc, pyspark.SparkContext)

    def test_set_spark_context_with_conf(self):
        """
        Check that a `SparkContext` object is correctly created for a given
        `SparkConf` object in the config dictionary.
        """
        conf = {"spark.app.name": "my-pyspark-app1"}
        sconf = pyspark.SparkConf().setAll(conf.items())
        sc = pyspark.SparkContext(conf=sconf)

        backend = Backend.SparkBackend(sparkcontext=sc)

        self.assertIsInstance(backend.sc, pyspark.SparkContext)
        appname = backend.sc.getConf().get("spark.app.name")
        self.assertEqual(appname, "my-pyspark-app1")

    def test_optimize_npartitions_with_num_executors(self):
        """
        Check that the number of partitions is correctly set to number of
        executors in the SparkConf dictionary.
        """
        conf = {"spark.executor.instances": 10}
        sconf = pyspark.SparkConf().setAll(conf.items())
        sc = pyspark.SparkContext(conf=sconf)
        backend = Backend.SparkBackend(sparkcontext=sc)

        self.assertEqual(backend.optimize_npartitions(), 10)

class OperationSupportTest(unittest.TestCase):
    """
    Ensure that incoming operations are classified accurately in distributed
    environment.
    """

    @classmethod
    def setUpClass(cls):
        """
        Synchronize PYSPARK_PYTHON variable to the current Python executable.

        Needed to avoid mismatch between python versions on driver and on
        the fake executor on the same machine.
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    @classmethod
    def tearDownClass(cls):
        """
        Stop the SparkContext and reset environment variable.
        """
        os.environ["PYSPARK_PYTHON"] = ""

    def test_action(self):
        """Check that action nodes are classified accurately."""
        backend = Backend.SparkBackend()
        backend.check_supported("Histo1D")

    def test_transformation(self):
        """Check that transformation nodes are classified accurately."""
        backend = Backend.SparkBackend()
        backend.check_supported("Define")

    def test_unsupported_operations(self):
        """Check that unsupported operations raise an Exception."""
        backend = Backend.SparkBackend()
        with self.assertRaises(Exception):
            backend.check_supported("Take")

        with self.assertRaises(Exception):
            backend.check_supported("Foreach")

        with self.assertRaises(Exception):
            backend.check_supported("Range")

    def test_none(self):
        """Check that incorrect operations raise an Exception."""
        backend = Backend.SparkBackend()
        with self.assertRaises(Exception):
            backend.check_supported("random")

    def test_range_operation_single_thread(self):
        """
        Check that 'Range' operation works in single-threaded mode and raises an
        Exception in multi-threaded mode.
        """
        backend = Backend.SparkBackend()
        with self.assertRaises(Exception):
            backend.check_supported("Range")

class InitializationTest(unittest.TestCase):
    """Check initialization method in the Spark backend"""

    @classmethod
    def setUpClass(cls):
        """
        Synchronize PYSPARK_PYTHON variable to the current Python executable.

        Needed to avoid mismatch between python versions on driver and on
        the fake executor on the same machine.
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    @classmethod
    def tearDownClass(cls):
        """
        Stop the SparkContext and reset environment variable.
        """
        os.environ["PYSPARK_PYTHON"] = ""

    def test_initialization(self):
        """
        Check that the user initialization method is assigned to the current
        backend.
        """
        def returnNumber(n):
            return n

        DistRDF.initialize(returnNumber, 123)

        # Dummy df just to retrieve the initialization function
        df = Spark.RDataFrame(10)
        f = df._headnode.backend.initialization

        self.assertEqual(f(), 123)

    def test_initialization_method(self):
        """
        Check initialization method in Spark backend.
        Define a method in the ROOT interpreter called getValue which returns
        the value defined by the user on the python side.
        """
        def init(value):
            import ROOT
            cpp_code = '''int userValue = %s ;''' % value
            ROOT.gInterpreter.ProcessLine(cpp_code)

        DistRDF.initialize(init, 123)
        # Spark backend has a limited list of supported methods, so we use
        # Histo1D which is a supported action.
        # The code below creates an RDataFrame instance with one single entry
        # and defines a column 'u' whose value is taken from the variable
        # 'userValue'.
        # This variable is only declared inside the ROOT interpreter, however
        # the value of the variable is passed by the user from the python side.
        # If the init function defined by the user is properly propagated to the
        # Spark backend, each workers will run the init function as a first step
        # and hence the variable 'userValue' will be defined at runtime.
        # As a result the define operation should read the variable 'userValue'
        # and assign it to the entries of the column 'u' (only one entry).
        # Finally, Histo1D returns a histogram filled with one value. The mean
        # of this single value has to be the value itself, independently of
        # the number of spawned workers.
        df = Spark.RDataFrame(1).Define(
            "u", "userValue").Histo1D("u")
        h = df.GetValue()
        self.assertEqual(h.GetMean(), 123)

class EmptyTreeErrorTest(unittest.TestCase):
    """
    Distributed execution fails when the tree has no entries.
    """

    @classmethod
    def setUpClass(cls):
        """
        Synchronize PYSPARK_PYTHON variable to the current Python executable.

        Needed to avoid mismatch between python versions on driver and on
        the fake executor on the same machine.
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    @classmethod
    def tearDownClass(cls):
        """
        Stop the SparkContext and reset environment variable.
        """
        os.environ["PYSPARK_PYTHON"] = ""

    def test_histo_from_empty_root_file(self):
        """
        Check that when performing operations with the distributed backend on
        an RDataFrame without entries, DistRDF raises an error.
        """

        # Create an RDataFrame from a file with an empty tree
        rdf = Spark.RDataFrame(
            "NOMINAL", "emptytree.root")
        histo = rdf.Histo1D(("empty", "empty", 10, 0, 10), "mybranch")

        # Get entries in the histogram, raises error
        with self.assertRaises(RuntimeError):
            histo.GetEntries()

class ChangeAttributeTest(unittest.TestCase):
    """Tests that check correct changes in the class attributes"""

    @classmethod
    def setUpClass(cls):
        """
        Synchronize PYSPARK_PYTHON variable to the current Python executable.

        Needed to avoid mismatch between python versions on driver and on
        the fake executor on the same machine.
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    @classmethod
    def tearDownClass(cls):
        """
        Stop the SparkContext and reset environment variable.
        """
        os.environ["PYSPARK_PYTHON"] = ""

    def test_change_attribute_when_npartitions_greater_than_clusters(self):
        """
        Check that the `npartitions class attribute is changed when it is
        greater than the number of clusters in the ROOT file.
        """

        treename = "TotemNtuple"
        filelist = ["Slimmed_ntuple.root"]
        df = Spark.RDataFrame(treename, filelist, npartitions=10)

        self.assertEqual(df._headnode.npartitions, 10)
        histo = df.Histo1D("track_rp_3.x")
        nentries = histo.GetEntries()

        self.assertEqual(nentries, 10)
        self.assertEqual(df._headnode.npartitions, 1)

    def test_optimize_npartitions_with_spark_config_options(self):
        """
        Check that relevant spark configuration options optimize the number of
        partitions.
        """

        conf = {"spark.executor.cores": 4, "spark.executor.instances": 4}
        sconf = pyspark.SparkConf().setAll(conf.items())
        scontext = pyspark.SparkContext(conf=sconf)

        df = Spark.RDataFrame(100, sparkcontext=scontext)

        # The number of partitions was optimized to be equal to
        # spark.executor.cores * spark.executor.instances
        self.assertEqual(df._headnode.npartitions, 16)

    def test_user_supplied_npartitions_have_precedence(self):
        """
        Check that even if spark configuration options could optimize the number
        of partitions, a user supplied value for npartitions takes precedence.
        """

        conf = {"spark.executor.cores": 4, "spark.executor.instances": 4}
        sconf = pyspark.SparkConf().setAll(conf.items())
        scontext = pyspark.SparkContext(conf=sconf)

        df = Spark.RDataFrame(100, sparkcontext=scontext, npartitions=4)

        # The number of partitions was supplied by the user.
        self.assertEqual(df._headnode.npartitions, 4)

if __name__ == "__main__":
    unittest.main()
<commit_msg>[DF] Suppress spurious ResourceWarning in distrdf-pyspark tests<commit_after>import os
import sys
import unittest
import warnings

import DistRDF
import pyspark
from DistRDF.Backends import Spark
from DistRDF.Backends.Spark import Backend

class SparkBackendInitTest(unittest.TestCase):
    """
    Tests to ensure that the instance variables
    of `Spark` class are set according to the
    input `config` dict.
    """

    @classmethod
    def setUpClass(cls):
        """
        Set up test environment for this class. Currently this includes:

        - Synchronize PYSPARK_PYTHON variable to the current Python executable.
          Needed to avoid mismatch between python versions on driver and on the
          fake executor on the same machine.
        - Ignore `ResourceWarning: unclosed socket` warning triggered by Spark.
          this is ignored by default in any application, but Python's unittest
          library overrides the default warning filters thus exposing this
          warning
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

        if sys.version_info.major >= 3:
            warnings.simplefilter("ignore", ResourceWarning)

    @classmethod
    def tearDownClass(cls):
        """Reset test environment."""
        os.environ["PYSPARK_PYTHON"] = ""

        if sys.version_info.major >= 3:
            warnings.simplefilter("default", ResourceWarning)

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    def test_set_spark_context_default(self):
        """
        Check that a `SparkContext` object is created with default options for
        the current system.
        """
        backend = Backend.SparkBackend()

        self.assertIsInstance(backend.sc, pyspark.SparkContext)

    def test_set_spark_context_with_conf(self):
        """
        Check that a `SparkContext` object is correctly created for a given
        `SparkConf` object in the config dictionary.
        """
        conf = {"spark.app.name": "my-pyspark-app1"}
        sconf = pyspark.SparkConf().setAll(conf.items())
        sc = pyspark.SparkContext(conf=sconf)

        backend = Backend.SparkBackend(sparkcontext=sc)

        self.assertIsInstance(backend.sc, pyspark.SparkContext)
        appname = backend.sc.getConf().get("spark.app.name")
        self.assertEqual(appname, "my-pyspark-app1")

    def test_optimize_npartitions_with_num_executors(self):
        """
        Check that the number of partitions is correctly set to number of
        executors in the SparkConf dictionary.
        """
        conf = {"spark.executor.instances": 10}
        sconf = pyspark.SparkConf().setAll(conf.items())
        sc = pyspark.SparkContext(conf=sconf)
        backend = Backend.SparkBackend(sparkcontext=sc)

        self.assertEqual(backend.optimize_npartitions(), 10)

class OperationSupportTest(unittest.TestCase):
    """
    Ensure that incoming operations are classified accurately in distributed
    environment.
    """

    @classmethod
    def setUpClass(cls):
        """
        Set up test environment for this class. Currently this includes:

        - Synchronize PYSPARK_PYTHON variable to the current Python executable.
          Needed to avoid mismatch between python versions on driver and on the
          fake executor on the same machine.
        - Ignore `ResourceWarning: unclosed socket` warning triggered by Spark.
          this is ignored by default in any application, but Python's unittest
          library overrides the default warning filters thus exposing this
          warning
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

        if sys.version_info.major >= 3:
            warnings.simplefilter("ignore", ResourceWarning)

    @classmethod
    def tearDownClass(cls):
        """Reset test environment."""
        os.environ["PYSPARK_PYTHON"] = ""

        if sys.version_info.major >= 3:
            warnings.simplefilter("default", ResourceWarning)

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    def test_action(self):
        """Check that action nodes are classified accurately."""
        backend = Backend.SparkBackend()
        backend.check_supported("Histo1D")

    def test_transformation(self):
        """Check that transformation nodes are classified accurately."""
        backend = Backend.SparkBackend()
        backend.check_supported("Define")

    def test_unsupported_operations(self):
        """Check that unsupported operations raise an Exception."""
        backend = Backend.SparkBackend()
        with self.assertRaises(Exception):
            backend.check_supported("Take")

        with self.assertRaises(Exception):
            backend.check_supported("Foreach")

        with self.assertRaises(Exception):
            backend.check_supported("Range")

    def test_none(self):
        """Check that incorrect operations raise an Exception."""
        backend = Backend.SparkBackend()
        with self.assertRaises(Exception):
            backend.check_supported("random")

    def test_range_operation_single_thread(self):
        """
        Check that 'Range' operation works in single-threaded mode and raises an
        Exception in multi-threaded mode.
        """
        backend = Backend.SparkBackend()
        with self.assertRaises(Exception):
            backend.check_supported("Range")

class InitializationTest(unittest.TestCase):
    """Check initialization method in the Spark backend"""

    @classmethod
    def setUpClass(cls):
        """
        Set up test environment for this class. Currently this includes:

        - Synchronize PYSPARK_PYTHON variable to the current Python executable.
          Needed to avoid mismatch between python versions on driver and on the
          fake executor on the same machine.
        - Ignore `ResourceWarning: unclosed socket` warning triggered by Spark.
          this is ignored by default in any application, but Python's unittest
          library overrides the default warning filters thus exposing this
          warning
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

        if sys.version_info.major >= 3:
            warnings.simplefilter("ignore", ResourceWarning)

    @classmethod
    def tearDownClass(cls):
        """Reset test environment."""
        os.environ["PYSPARK_PYTHON"] = ""

        if sys.version_info.major >= 3:
            warnings.simplefilter("default", ResourceWarning)

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    def test_initialization(self):
        """
        Check that the user initialization method is assigned to the current
        backend.
        """
        def returnNumber(n):
            return n

        DistRDF.initialize(returnNumber, 123)

        # Dummy df just to retrieve the initialization function
        df = Spark.RDataFrame(10)
        f = df._headnode.backend.initialization

        self.assertEqual(f(), 123)

    def test_initialization_method(self):
        """
        Check initialization method in Spark backend.
        Define a method in the ROOT interpreter called getValue which returns
        the value defined by the user on the python side.
        """
        def init(value):
            import ROOT
            cpp_code = '''int userValue = %s ;''' % value
            ROOT.gInterpreter.ProcessLine(cpp_code)

        DistRDF.initialize(init, 123)
        # Spark backend has a limited list of supported methods, so we use
        # Histo1D which is a supported action.
        # The code below creates an RDataFrame instance with one single entry
        # and defines a column 'u' whose value is taken from the variable
        # 'userValue'.
        # This variable is only declared inside the ROOT interpreter, however
        # the value of the variable is passed by the user from the python side.
        # If the init function defined by the user is properly propagated to the
        # Spark backend, each workers will run the init function as a first step
        # and hence the variable 'userValue' will be defined at runtime.
        # As a result the define operation should read the variable 'userValue'
        # and assign it to the entries of the column 'u' (only one entry).
        # Finally, Histo1D returns a histogram filled with one value. The mean
        # of this single value has to be the value itself, independently of
        # the number of spawned workers.
        df = Spark.RDataFrame(1).Define(
            "u", "userValue").Histo1D("u")
        h = df.GetValue()
        self.assertEqual(h.GetMean(), 123)

class EmptyTreeErrorTest(unittest.TestCase):
    """
    Distributed execution fails when the tree has no entries.
    """

    @classmethod
    def setUpClass(cls):
        """
        Set up test environment for this class. Currently this includes:

        - Synchronize PYSPARK_PYTHON variable to the current Python executable.
          Needed to avoid mismatch between python versions on driver and on the
          fake executor on the same machine.
        - Ignore `ResourceWarning: unclosed socket` warning triggered by Spark.
          this is ignored by default in any application, but Python's unittest
          library overrides the default warning filters thus exposing this
          warning
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

        if sys.version_info.major >= 3:
            warnings.simplefilter("ignore", ResourceWarning)

    @classmethod
    def tearDownClass(cls):
        """Reset test environment."""
        os.environ["PYSPARK_PYTHON"] = ""

        if sys.version_info.major >= 3:
            warnings.simplefilter("default", ResourceWarning)

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    def test_histo_from_empty_root_file(self):
        """
        Check that when performing operations with the distributed backend on
        an RDataFrame without entries, DistRDF raises an error.
        """

        # Create an RDataFrame from a file with an empty tree
        rdf = Spark.RDataFrame(
            "NOMINAL", "emptytree.root")
        histo = rdf.Histo1D(("empty", "empty", 10, 0, 10), "mybranch")

        # Get entries in the histogram, raises error
        with self.assertRaises(RuntimeError):
            histo.GetEntries()

class ChangeAttributeTest(unittest.TestCase):
    """Tests that check correct changes in the class attributes"""

    @classmethod
    def setUpClass(cls):
        """
        Set up test environment for this class. Currently this includes:

        - Synchronize PYSPARK_PYTHON variable to the current Python executable.
          Needed to avoid mismatch between python versions on driver and on the
          fake executor on the same machine.
        - Ignore `ResourceWarning: unclosed socket` warning triggered by Spark.
          this is ignored by default in any application, but Python's unittest
          library overrides the default warning filters thus exposing this
          warning
        """
        os.environ["PYSPARK_PYTHON"] = sys.executable

        if sys.version_info.major >= 3:
            warnings.simplefilter("ignore", ResourceWarning)

    @classmethod
    def tearDownClass(cls):
        """Reset test environment."""
        os.environ["PYSPARK_PYTHON"] = ""

        if sys.version_info.major >= 3:
            warnings.simplefilter("default", ResourceWarning)

    def tearDown(self):
        """Stop any created SparkContext"""
        pyspark.SparkContext.getOrCreate().stop()

    def test_change_attribute_when_npartitions_greater_than_clusters(self):
        """
        Check that the `npartitions class attribute is changed when it is
        greater than the number of clusters in the ROOT file.
        """

        treename = "TotemNtuple"
        filelist = ["Slimmed_ntuple.root"]
        df = Spark.RDataFrame(treename, filelist, npartitions=10)

        self.assertEqual(df._headnode.npartitions, 10)
        histo = df.Histo1D("track_rp_3.x")
        nentries = histo.GetEntries()

        self.assertEqual(nentries, 10)
        self.assertEqual(df._headnode.npartitions, 1)

    def test_optimize_npartitions_with_spark_config_options(self):
        """
        Check that relevant spark configuration options optimize the number of
        partitions.
        """

        conf = {"spark.executor.cores": 4, "spark.executor.instances": 4}
        sconf = pyspark.SparkConf().setAll(conf.items())
        scontext = pyspark.SparkContext(conf=sconf)

        df = Spark.RDataFrame(100, sparkcontext=scontext)

        # The number of partitions was optimized to be equal to
        # spark.executor.cores * spark.executor.instances
        self.assertEqual(df._headnode.npartitions, 16)

    def test_user_supplied_npartitions_have_precedence(self):
        """
        Check that even if spark configuration options could optimize the number
        of partitions, a user supplied value for npartitions takes precedence.
        """

        conf = {"spark.executor.cores": 4, "spark.executor.instances": 4}
        sconf = pyspark.SparkConf().setAll(conf.items())
        scontext = pyspark.SparkContext(conf=sconf)

        df = Spark.RDataFrame(100, sparkcontext=scontext, npartitions=4)

        # The number of partitions was supplied by the user.
        self.assertEqual(df._headnode.npartitions, 4)

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before>#-*- coding: utf-8 -*-
__version__ = '3.1.3.2'
<commit_msg>Bump to 3.1.4.1<commit_after>#-*- coding: utf-8 -*-
__version__ = '3.1.4.1'
<|endoftext|>


<commit_before><commit_msg>Adding get_config_reply and set_config base class<commit_after>"""Defines Features Reply classes and related items"""

# System imports
import enum

# Third-party imports

# Local source tree imports
from common import base
from common import basic_types
from common_structures import header as of_header

# Enums

class ConfigFlags(enum.Enum):
    """Configuratin Flags

        # Handling of IP Fragments
        OFPC_FRAG_NORMAL               # No special handling for fragments
        OFPC_FRAG_DROP                 # Drop fragments
        OFPC_FRAG_REASM                # Reassemble (only if OFPC_IP_REASM set)
        OFPC_FRAG_MASK

        # TTL processing - applicable for IP and MPLS packets
        OFPC_INVALID_TTL_TO_CONTROLLER # Send packets with invalid TTL
                                       # i.e. 0 or 1 to controller
    """
    OFPC_FRAG_NORMAL = 0
    OFPC_FRAG_DROP = 1 << 0
    OFPC_FRAG_REASM = 1 << 1
    OFPC_FRAG_MASK = 3
    OFPC_INVALID_TTL_TO_CONTROLLER = 1 << 2

# Classes

class SwitchConig(base.GenericStruct):
    """Message used on OFPT_SET_CONFIG and OFPT_GET_CONFIG_REPLY.

    This is OFPT_GET_CONFIG_REPLY message sent by the switch to
    the controller in response to the OFPT_GET_CONFIG_REQUEST,
    and also is the OFPT_SET_CONFIG message sent by the controller
    to the switch.

        :param xid:           xid to be used on the message header
        :param set_message:   Boolean to define if the message is
                              a SET_CONFIG or GET_CONFIG_REPLY.
                              - True (default): OFPT_SET_CONFIG
                              - False: OFPT_GET_CONFIG_REPLY
        :param flags:         UBInt16 OFPC_* flags
        :param miss_send_len: UBInt16 max bytes of new flow that the
                              datapath should send to the controller
    """

    header = of_header.OFPHeader()
    flags = basic_types.UBInt16
    miss_send_len = basic_types.UBInt16

    def __init__(self,
                 xid=None,
                 set_message=True,
                 flags=None,
                 miss_send_len=0):
        if set_message:
            self.header.ofp_type = of_header.OFPType.OFPT_SET_CONFIG
        else:
            self.header.ofp_type = of_header.OFPType.OFPT_GET_CONFIG_REPLY
        self.header.xid = xid
        self.flags = flags
        self.miss_send_len = miss_send_len
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright (C) 2011-2013 by Florian Mounier, Kozea
# This file is part of pystil, licensed under a 3-clause BSD license.

"""Utility functions to help processing data"""
from datetime import timedelta
from urllib.parse import urlparse, parse_qs
from pystil.context import pystil
from pystil.db import Visit
from sqlalchemy.sql import literal, func
import re

IPV4RE = re.compile(r"(\d{1,3}(\.|$)){4}")
BROWSER_VERSION_NUMBERS = {
    'opera': 1,
    'safari': 1,
    'chrome': 1}

# Stripped down user agent parsing from werkzeug
PLATFORMS = [(b, re.compile(a, re.I)) for a, b in (
    ('iphone|ios', 'iphone'),
    (r'darwin|mac|os\s*x', 'macos'),
    ('win', 'windows'),
    (r'android', 'android'),
    (r'x11|lin(\b|ux)?', 'linux'),
    ('(sun|i86)os', 'solaris'),
    (r'nintendo\s+wii', 'wii'),
    ('irix', 'irix'),
    ('hp-?ux', 'hpux'),
    ('aix', 'aix'),
    ('sco|unix_sv', 'sco'),
    ('bsd', 'bsd'),
    ('amiga', 'amiga')
)]

BROWSERS = [(
    b, re.compile(r'(?:%s)[/\sa-z(]*(\d+[.\da-z]+)?(?i)' % a)
) for a, b in (
    ('googlebot', 'google'),
    ('msnbot', 'msn'),
    ('yahoo', 'yahoo'),
    ('ask jeeves', 'ask'),
    (r'aol|america\s+online\s+browser', 'aol'),
    ('opera', 'opera'),
    ('chrome', 'chrome'),
    ('firefox|firebird|phoenix|iceweasel', 'firefox'),
    ('galeon', 'galeon'),
    ('safari', 'safari'),
    ('webkit', 'webkit'),
    ('camino', 'camino'),
    ('konqueror', 'konqueror'),
    ('k-meleon', 'kmeleon'),
    ('netscape', 'netscape'),
    (r'msie|microsoft\s+internet\s+explorer', 'msie'),
    ('lynx', 'lynx'),
    ('links', 'links'),
    ('seamonkey|mozilla', 'seamonkey')
)]

def on(host, table=Visit.__table__):
    """Generate a filter on a host """
    if host == 'all':
        return literal(True)
    return table.c.host.like('%' + host)

def between(from_date, to_date, table=Visit.__table__):
    """Generate a filter between 2 dates"""
    return ((from_date <= table.c.date) &
            (table.c.date < to_date + timedelta(1)))

def visit_to_table_line(visit):
    html = '
' % visit.uuid
    for key in ['date', 'site', 'ip', 'country', 'city', 'page', 'referrer']:
        html += ''
        val = getattr(visit, key)
        if val:
            if key == 'date':
                val = val.strftime('%Y-%m-%d %H:%M:%S')
            html += val
        html += '	'
    html += '
'
    return html

def parse_referrer(referrer, with_query=False, host_only=False,
                   second_pass=False):
    """Return a pretty format for most search engines"""
    if referrer:
        up = urlparse(referrer)
        netloc = up.netloc
        if not netloc:
            if with_query:
                return u"Local: %s" % up.path
            if second_pass:
                return referrer
            return "Local"

        query = up.query
        parsed = parse_qs(query)
        search = parsed.get('q', parsed.get('p', parsed.get('rdata', None)))
        if search:
            # TODO Yahoo variable encoding
            if with_query:
                return u"Organic: %s %s" % (netloc, try_decode(search[0]))
            return u"Organic: %s" % netloc
        if host_only:
            return netloc
        return referrer
    return "Direct"

def parse_domain(referrer):
    """Return the domain of the given referrer."""
    if referrer:
        up = urlparse(referrer)
        netloc = up.netloc
        return netloc or "Local"
    return "Direct"

def get_aggregate(criteria):
    """Returns a tuple in the form (model, count_column).
    Tries to look for an aggregate table suitable for the criteria
    """
    model = getattr(pystil.db, 'Agg_by_%s' % criteria, Visit)
    if model == Visit:
        pystil.log.warn('No aggregates for criteria %s' % criteria)
        count = func.count(1)
    else:
        count = func.sum('fact_count')
    return model, count

def try_decode(astring):
    """Try decoding a string in various encodings, with a fallback to good old
    ascii."""
    for encoding in ('utf8', 'latin'):
        try:
            return astring.decode(encoding)
        except UnicodeDecodeError:
            pass
    return astring.decode('ascii', 'ignore')

def parse_ua(user_agent):
    """Parse the user agent to return the platform, browser and version"""
    for platform, regex in PLATFORMS:
        match = regex.search(user_agent)
        if match is not None:
            break
    else:
        platform = None
    for browser, regex in BROWSERS:
        match = regex.search(user_agent)
        if match is not None:
            version = match.group(1)
            break
    else:
        browser = version = None
    return platform, browser, version
<commit_msg>Pretty referrers<commit_after># -*- coding: utf-8 -*-
# Copyright (C) 2011-2013 by Florian Mounier, Kozea
# This file is part of pystil, licensed under a 3-clause BSD license.

"""Utility functions to help processing data"""
from datetime import timedelta
from urllib.parse import urlparse, parse_qs
from pystil.context import pystil
from pystil.db import Visit
from sqlalchemy.sql import literal, func
import re

IPV4RE = re.compile(r"(\d{1,3}(\.|$)){4}")
BROWSER_VERSION_NUMBERS = {
    'opera': 1,
    'safari': 1,
    'chrome': 1}

# Stripped down user agent parsing from werkzeug
PLATFORMS = [(b, re.compile(a, re.I)) for a, b in (
    ('iphone|ios', 'iphone'),
    (r'darwin|mac|os\s*x', 'macos'),
    ('win', 'windows'),
    (r'android', 'android'),
    (r'x11|lin(\b|ux)?', 'linux'),
    ('(sun|i86)os', 'solaris'),
    (r'nintendo\s+wii', 'wii'),
    ('irix', 'irix'),
    ('hp-?ux', 'hpux'),
    ('aix', 'aix'),
    ('sco|unix_sv', 'sco'),
    ('bsd', 'bsd'),
    ('amiga', 'amiga')
)]

BROWSERS = [(
    b, re.compile(r'(?:%s)[/\sa-z(]*(\d+[.\da-z]+)?(?i)' % a)
) for a, b in (
    ('googlebot', 'google'),
    ('msnbot', 'msn'),
    ('yahoo', 'yahoo'),
    ('ask jeeves', 'ask'),
    (r'aol|america\s+online\s+browser', 'aol'),
    ('opera', 'opera'),
    ('chrome', 'chrome'),
    ('firefox|firebird|phoenix|iceweasel', 'firefox'),
    ('galeon', 'galeon'),
    ('safari', 'safari'),
    ('webkit', 'webkit'),
    ('camino', 'camino'),
    ('konqueror', 'konqueror'),
    ('k-meleon', 'kmeleon'),
    ('netscape', 'netscape'),
    (r'msie|microsoft\s+internet\s+explorer', 'msie'),
    ('lynx', 'lynx'),
    ('links', 'links'),
    ('seamonkey|mozilla', 'seamonkey')
)]

def on(host, table=Visit.__table__):
    """Generate a filter on a host """
    if host == 'all':
        return literal(True)
    return table.c.host.like('%' + host)

def between(from_date, to_date, table=Visit.__table__):
    """Generate a filter between 2 dates"""
    return ((from_date <= table.c.date) &
            (table.c.date < to_date + timedelta(1)))

def visit_to_table_line(visit):
    html = '
' % visit.uuid
    for key in [
            'date', 'site', 'ip', 'country',
            'city', 'page', 'pretty_referrer']:
        html += ''
        val = getattr(visit, key)
        if val:
            if key == 'date':
                val = val.strftime('%Y-%m-%d %H:%M:%S')
            html += val
        html += '	'
    html += '
'
    return html

def parse_referrer(referrer, with_query=False, host_only=False,
                   second_pass=False):
    """Return a pretty format for most search engines"""
    if referrer:
        up = urlparse(referrer)
        netloc = up.netloc
        if not netloc:
            if with_query:
                return u"Local: %s" % up.path
            if second_pass:
                return referrer
            return "Local"

        query = up.query
        parsed = parse_qs(query)
        search = parsed.get('q', parsed.get('p', parsed.get('rdata', None)))
        if search:
            # TODO Yahoo variable encoding
            if with_query:
                return u"Organic: %s %s" % (netloc, try_decode(search[0]))
            return u"Organic: %s" % netloc
        if host_only:
            return netloc
        return referrer
    return "Direct"

def parse_domain(referrer):
    """Return the domain of the given referrer."""
    if referrer:
        up = urlparse(referrer)
        netloc = up.netloc
        return netloc or "Local"
    return "Direct"

def get_aggregate(criteria):
    """Returns a tuple in the form (model, count_column).
    Tries to look for an aggregate table suitable for the criteria
    """
    model = getattr(pystil.db, 'Agg_by_%s' % criteria, Visit)
    if model == Visit:
        pystil.log.warn('No aggregates for criteria %s' % criteria)
        count = func.count(1)
    else:
        count = func.sum('fact_count')
    return model, count

def try_decode(astring):
    """Try decoding a string in various encodings, with a fallback to good old
    ascii."""
    for encoding in ('utf8', 'latin'):
        try:
            return astring.decode(encoding)
        except UnicodeDecodeError:
            pass
    return astring.decode('ascii', 'ignore')

def parse_ua(user_agent):
    """Parse the user agent to return the platform, browser and version"""
    for platform, regex in PLATFORMS:
        match = regex.search(user_agent)
        if match is not None:
            break
    else:
        platform = None
    for browser, regex in BROWSERS:
        match = regex.search(user_agent)
        if match is not None:
            version = match.group(1)
            break
    else:
        browser = version = None
    return platform, browser, version
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import unicode_literals

import StringIO

from datetime import datetime, timedelta

from django.conf import settings
from django.core.cache import cache as memcache
from django.core.validators import MaxValueValidator
from django.db import models
from django.utils.encoding import python_2_unicode_compatible
from django.utils.translation import ugettext_lazy as _

from cms.models.pluginmodel import CMSPlugin

from .utils import get_cache_key

from github import Github

CACHE_DURATION = getattr(settings, "ALDRYN_GITHUB_STATS_CACHE_DURATION", 3600)

@python_2_unicode_compatible
class GitHubStatsRepository(models.Model):

    label = models.CharField(
        max_length=128, default='', blank=False,
        help_text=_('Provide a descriptive label for your repo. E.g., '
                    '"django CMS'))
    full_name = models.CharField(
        max_length=255, blank=False, default='',
        help_text=_('Enter the repo. full name. E.g., "divio/django-cms"'),)
    token = models.CharField(
        max_length=64, blank=False, default='',
        help_text=_('Provide a suitable GitHub API token.'))

    class Meta:
        verbose_name = _('repository')
        verbose_name_plural = _('repositories')
        unique_together = ('label', 'full_name', 'token')

    def __str__(self):
        return self.label

    def get_cache_key(self):
        """
        Gets the cache key for this specific repo configuration.
        """
        return get_cache_key(
            self.__class__.__name__, settings=(self.pk, ))

    def get_repo(self, force_refresh=False):
        """
        Using the configuration for this GitHubStatsRepository instance, fetch
        the repo object from cache, if possible, else, from GitHub.

        :param force_refresh: Set to force an update from GitHub.
        :type force_refresh: bool
        :rtype: GitHubStatsRepository
        """
        key = self.get_cache_key()
        g = Github(self.token)
        if force_refresh:
            repo = None
        else:
            pickled = memcache.get(key, None)
            if pickled:
                file_str = StringIO.StringIO(pickled)
                repo = g.load(file_str)
            else:
                repo = None
        if repo is None:
            repo = g.get_repo(self.full_name)
            output = StringIO.StringIO()
            g.dump(repo, output)
            duration = CACHE_DURATION * 2 if force_refresh else CACHE_DURATION
            memcache.set(key, output.getvalue(), duration)
        return repo

class GitHubStatsBase(CMSPlugin):
    # Avoid reverse relation name clashes by not adding a related_name
    # to the parent plugin
    cmsplugin_ptr = models.OneToOneField(
        CMSPlugin, related_name='+', parent_link=True)
    repo = models.ForeignKey(
        'GitHubStatsRepository', null=True, verbose_name=_('repository'),
        help_text=_('Select the repository to work with.'))
    subhead = models.CharField(
        _('subhead'), max_length=255, blank=True, default='',
        help_text=_('Optional subheading.'))
    subhead_link = models.URLField(
        _('URL'), max_length=4096, blank=True, default='',
        help_text=_('Optional subhead link destination.'))

    class Meta:
        abstract = True

@python_2_unicode_compatible
class GitHubStatsRecentCommitsPluginModel(GitHubStatsBase):

    from_days_ago = models.PositiveIntegerField(
        _('Number of days ago'), default=30,
        validators=[MaxValueValidator(365), ],
        help_text=_('Number of days to sum commits (maximum 365)'))

    def recent_commits(self):
        if not self.repo or not self.repo.full_name or not self.repo.token:
            return 0
        key = get_cache_key(
            self.__class__.__name__,
            settings=(self.repo.full_name, self.repo.token, self.from_days_ago))

        cached_value = memcache.get(key)
        if cached_value is None:
            repo = self.repo.get_repo()
            if repo:
                today = datetime.today()
                days_ago = today - timedelta(days=self.from_days_ago or 0)
                total = 0
                for commits in repo.get_stats_commit_activity():
                    delta = commits.week - days_ago
                    if delta.days > -7:
                        if delta.days < 0:
                            for num in commits.days[-delta.days:]:
                                total += num
                        else:
                            total += commits.total
                cached_value = total
                if total is not None:
                    memcache.set(key, total, CACHE_DURATION)
            else:
                return 0
        return cached_value

    def __str__(self):
        return 'Recent commits for %d days on %s' % (
            self.from_days_ago,
            self.repo.full_name if self.repo else '',
        )

@python_2_unicode_compatible
class GitHubStatsIssuesCountPluginModel(GitHubStatsBase):

    CHOICES = (
        ('open', _('Open'), ),
        ('closed', _('Closed'), ),
    )

    from_days_ago = models.PositiveIntegerField(
        _('Number of days ago'), default=30,
        validators=[MaxValueValidator(365), ],
        help_text=_('Number of days to sum commits (maximum 365)'))
    state = models.CharField(
        max_length=16, choices=CHOICES, blank=False, default=CHOICES[1][0],
        help_text=_('Event type'))

    def recent_issues(self):
        if not self.repo or not self.repo.full_name or not self.repo.token:
            return 0
        key = get_cache_key(
            self.__class__.__name__,
            settings=(self.repo.full_name, self.repo.token, self.from_days_ago,
                      self.state))
        cached_value = memcache.get(key)
        if cached_value is None:
            repo = self.repo.get_repo()
            if repo:
                today = datetime.today()
                days_ago = today - timedelta(days=self.from_days_ago)

                issues = repo.get_issues(
                    state=self.state, since=days_ago)
                cached_value = len(list(issues))
                if cached_value is not None:
                    memcache.set(key, cached_value, CACHE_DURATION)
            else:
                return 0
        return cached_value

    def __str__(self):
        return 'Recent %s issues for %d days on %s' % (
            self.state.lower(),
            self.from_days_ago,
            self.repo.full_name if self.repo else ''
        )

@python_2_unicode_compatible
class GitHubStatsRepoPropertyPluginModel(GitHubStatsBase):

    CHOICES = (
        ('forks_count', _('No. forks'), ),
        ('network_count', _('No. networks'), ),
        ('open_issues_count', _('No. open issues'), ),
        ('size', _('Size'), ),
        ('stargazers_count', _('No. stargazers'), ),
        ('watchers_count', _('No. watchers'), ),
    )

    property_name = models.CharField(_('property'), max_length=64,
        blank=False, default=CHOICES[0][0], choices=CHOICES,
        help_text=_('Choose a repository property to display.'))

    property_label = models.CharField(_('label'), max_length=128,
        blank=False, default='',
        help_text=_('Label to display.'))

    def property(self):
        if not self.repo or not self.repo.full_name or not self.repo.token:
            return 0
        key = get_cache_key(
            self.__class__.__name__,
            settings=(self.repo.full_name, self.repo.token, self.property_name))
        cached_value = memcache.get(key)
        if cached_value is None:
            repo = self.repo.get_repo()
            if repo:
                cached_value = getattr(repo, self.property_name)
                if cached_value is not None:
                    memcache.set(key, cached_value, CACHE_DURATION)
            else:
                return 0
        return cached_value

    def __str__(self):
        return 'Property %s of "%s"' % (
            self.property_name,
            self.repo.full_name if self.repo else 'unset repo',
        )
<commit_msg>Adds logging, renames var to 'cache'<commit_after># -*- coding: utf-8 -*-

from __future__ import unicode_literals

import logging
import StringIO

from datetime import datetime, timedelta

from django.conf import settings
from django.core.cache import cache
from django.core.validators import MaxValueValidator
from django.db import models
from django.utils.encoding import python_2_unicode_compatible
from django.utils.translation import ugettext_lazy as _

from cms.models.pluginmodel import CMSPlugin

from .utils import get_cache_key

from github import Github

logger = logging.getLogger(__name__)

CACHE_DURATION = getattr(settings, "ALDRYN_GITHUB_STATS_CACHE_DURATION", 3600)

@python_2_unicode_compatible
class GitHubStatsRepository(models.Model):

    label = models.CharField(
        max_length=128, default='', blank=False,
        help_text=_('Provide a descriptive label for your repo. E.g., '
                    '"django CMS'))
    full_name = models.CharField(
        max_length=255, blank=False, default='',
        help_text=_('Enter the repo. full name. E.g., "divio/django-cms"'),)
    token = models.CharField(
        max_length=64, blank=False, default='',
        help_text=_('Provide a suitable GitHub API token.'))

    class Meta:
        verbose_name = _('repository')
        verbose_name_plural = _('repositories')
        unique_together = ('label', 'full_name', 'token')

    def __str__(self):
        return self.label

    def get_cache_key(self):
        """
        Gets the cache key for this specific repo configuration.
        """
        return get_cache_key(
            self.__class__.__name__, settings=(self.pk, ))

    def get_repo(self, force_refresh=False):
        """
        Using the configuration for this GitHubStatsRepository instance, fetch
        the repo object from cache, if possible, else, from GitHub.

        :param force_refresh: Set to force an update from GitHub.
        :type force_refresh: bool
        :rtype: GitHubStatsRepository
        """
        key = self.get_cache_key()
        g = Github(self.token)
        if force_refresh:
            repo = None
        else:
            pickled = cache.get(key, None)
            if pickled:
                file_str = StringIO.StringIO(pickled)
                repo = g.load(file_str)
            else:
                repo = None
        if repo is None:
            if force_refresh:
                logger.info('Forced refresh')
            else:
                logger.info('Natural refresh')
            repo = g.get_repo(self.full_name)
            output = StringIO.StringIO()
            g.dump(repo, output)
            duration = CACHE_DURATION * 2 if force_refresh else CACHE_DURATION
            cache.set(key, output.getvalue(), duration)
        return repo

class GitHubStatsBase(CMSPlugin):
    # Avoid reverse relation name clashes by not adding a related_name
    # to the parent plugin
    cmsplugin_ptr = models.OneToOneField(
        CMSPlugin, related_name='+', parent_link=True)
    repo = models.ForeignKey(
        'GitHubStatsRepository', null=True, verbose_name=_('repository'),
        help_text=_('Select the repository to work with.'))
    subhead = models.CharField(
        _('subhead'), max_length=255, blank=True, default='',
        help_text=_('Optional subheading.'))
    subhead_link = models.URLField(
        _('URL'), max_length=4096, blank=True, default='',
        help_text=_('Optional subhead link destination.'))

    class Meta:
        abstract = True

@python_2_unicode_compatible
class GitHubStatsRecentCommitsPluginModel(GitHubStatsBase):

    from_days_ago = models.PositiveIntegerField(
        _('Number of days ago'), default=30,
        validators=[MaxValueValidator(365), ],
        help_text=_('Number of days to sum commits (maximum 365)'))

    def recent_commits(self):
        if not self.repo or not self.repo.full_name or not self.repo.token:
            return 0
        key = get_cache_key(
            self.__class__.__name__,
            settings=(self.repo.full_name, self.repo.token, self.from_days_ago))

        cached_value = cache.get(key)
        if cached_value is None:
            repo = self.repo.get_repo()
            if repo:
                today = datetime.today()
                days_ago = today - timedelta(days=self.from_days_ago or 0)
                total = 0
                for commits in repo.get_stats_commit_activity():
                    delta = commits.week - days_ago
                    if delta.days > -7:
                        if delta.days < 0:
                            for num in commits.days[-delta.days:]:
                                total += num
                        else:
                            total += commits.total
                cached_value = total
                if total is not None:
                    cache.set(key, total, CACHE_DURATION)
            else:
                return 0
        return cached_value

    def __str__(self):
        return 'Recent commits for %d days on %s' % (
            self.from_days_ago,
            self.repo.full_name if self.repo else '',
        )

@python_2_unicode_compatible
class GitHubStatsIssuesCountPluginModel(GitHubStatsBase):

    CHOICES = (
        ('open', _('Open'), ),
        ('closed', _('Closed'), ),
    )

    from_days_ago = models.PositiveIntegerField(
        _('Number of days ago'), default=30,
        validators=[MaxValueValidator(365), ],
        help_text=_('Number of days to sum commits (maximum 365)'))
    state = models.CharField(
        max_length=16, choices=CHOICES, blank=False, default=CHOICES[1][0],
        help_text=_('Event type'))

    def recent_issues(self):
        if not self.repo or not self.repo.full_name or not self.repo.token:
            return 0
        key = get_cache_key(
            self.__class__.__name__,
            settings=(self.repo.full_name, self.repo.token, self.from_days_ago,
                      self.state))
        cached_value = cache.get(key)
        if cached_value is None:
            repo = self.repo.get_repo()
            if repo:
                today = datetime.today()
                days_ago = today - timedelta(days=self.from_days_ago)

                issues = repo.get_issues(
                    state=self.state, since=days_ago)
                cached_value = len(list(issues))
                if cached_value is not None:
                    cache.set(key, cached_value, CACHE_DURATION)
            else:
                return 0
        return cached_value

    def __str__(self):
        return 'Recent %s issues for %d days on %s' % (
            self.state.lower(),
            self.from_days_ago,
            self.repo.full_name if self.repo else ''
        )

@python_2_unicode_compatible
class GitHubStatsRepoPropertyPluginModel(GitHubStatsBase):

    CHOICES = (
        ('forks_count', _('No. forks'), ),
        ('network_count', _('No. networks'), ),
        ('open_issues_count', _('No. open issues'), ),
        ('size', _('Size'), ),
        ('stargazers_count', _('No. stargazers'), ),
        ('watchers_count', _('No. watchers'), ),
    )

    property_name = models.CharField(_('property'), max_length=64,
        blank=False, default=CHOICES[0][0], choices=CHOICES,
        help_text=_('Choose a repository property to display.'))

    property_label = models.CharField(_('label'), max_length=128,
        blank=False, default='',
        help_text=_('Label to display.'))

    def property(self):
        if not self.repo or not self.repo.full_name or not self.repo.token:
            return 0
        key = get_cache_key(
            self.__class__.__name__,
            settings=(self.repo.full_name, self.repo.token, self.property_name))
        cached_value = cache.get(key)
        if cached_value is None:
            repo = self.repo.get_repo()
            if repo:
                cached_value = getattr(repo, self.property_name)
                if cached_value is not None:
                    cache.set(key, cached_value, CACHE_DURATION)
            else:
                return 0
        return cached_value

    def __str__(self):
        return 'Property %s of "%s"' % (
            self.property_name,
            self.repo.full_name if self.repo else 'unset repo',
        )
<|endoftext|>


<commit_before>import logging

from dateutil.parser import parse
from .utils import ApiComponent

log = logging.getLogger(__name__)

class Team(ApiComponent):
    """ A Microsoft Teams team """

    _endpoints = {}

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Microsoft Teams team

        :param parent: parent object
        :type parent: Teams
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        cloud_data = kwargs.get(self._cloud_data_key, {})

        self.object_id = cloud_data.get('id')

        # Choose the main_resource passed in kwargs over parent main_resource
        main_resource = kwargs.pop('main_resource', None) or (
            getattr(parent, 'main_resource', None) if parent else None)

        main_resource = '{}{}'.format(main_resource, '')

        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

        self.display_name = cloud_data.get(self._cc('displayName'), '')
        self.description = cloud_data.get(self._cc('description'), '')
        self.is_archived = cloud_data.get(self._cc('isArchived'), '')
        self.web_url = cloud_data.get(self._cc('webUrl'), '')

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'Team: {}'.format(self.displayName)

    def __eq__(self, other):
        return self.object_id == other.object_id

class Channel(ApiComponent):
    """ A Microsoft Teams channel """

    _endpoints = {}

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Microsoft Teams channel

        :param parent: parent object
        :type parent: Teams
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        cloud_data = kwargs.get(self._cloud_data_key, {})

        self.object_id = cloud_data.get('id')

        # Choose the main_resource passed in kwargs over parent main_resource
        main_resource = kwargs.pop('main_resource', None) or (
            getattr(parent, 'main_resource', None) if parent else None)

        main_resource = '{}{}'.format(main_resource, '')

        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

        self.display_name = cloud_data.get(self._cc('displayName'), '')
        self.description = cloud_data.get('description')
        self.email = cloud_data.get('email')

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'Channel: {}'.format(self.displayName)

    def __eq__(self, other):
        return self.object_id == other.object_id

class App(ApiComponent):
    """ A Microsoft Teams app """

    _endpoints = {}

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Microsoft Teams app

        :param parent: parent object
        :type parent: Teams
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        cloud_data = kwargs.get(self._cloud_data_key, {})

        self.object_id = cloud_data.get('id')

        # Choose the main_resource passed in kwargs over parent main_resource
        main_resource = kwargs.pop('main_resource', None) or (
            getattr(parent, 'main_resource', None) if parent else None)

        main_resource = '{}{}'.format(main_resource, '')

        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

        self.app_definition = cloud_data.get(self._cc('teamsAppDefinition'), '')

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'App: {}'.format(self.app_definition['displayName'])

    def __eq__(self, other):
        return self.object_id == other.object_id

class Teams(ApiComponent):
    """ A microsoft teams class
        In order to use the API following permissions are required.
        Delegated (work or school account) - Group.Read.All, Group.ReadWrite.All
    """

    _endpoints = {
        'get_my_teams': '/me/joinedTeams',
        'get_channels': '/teams/{team_id}/channels',
        'create_channel': '/teams/{team_id}/channels',
        'get_channel_info': '/teams/{team_id}/channels/{channel_id}',
        'get_apps_in_team': '/teams/{team_id}/installedApps?$expand=teamsAppDefinition',
    }

    team_constructor = Team
    channel_constructor = Channel
    app_constructor = App

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Teams object

        :param parent: parent object
        :type parent: Account
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        # Choose the main_resource passed in kwargs over the host_name
        main_resource = kwargs.pop('main_resource',
                                   '')  # defaults to blank resource
        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'Microsoft Teams'

    def get_my_teams(self, *args):
        """ Returns a list of teams that I am in

        :rtype: teams
        """

        url = self.build_url(self._endpoints.get('get_my_teams'))

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return [
            self.team_constructor(parent=self, **{self._cloud_data_key: site})
            for site in data.get('value', [])]

    def get_channels(self, team_id=None):
        """ Returns a list of channels of a specified team

        :param team_id: the team_id of the channel to be retrieved.

        :rtype: channels
        """

        if not team_id:
            raise RuntimeError('Provide the team_id')

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_channels').format(team_id=team_id))

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return [
            self.channel_constructor(parent=self, **{self._cloud_data_key: site})
            for site in data.get('value', [])]

    def create_channel(self, team_id=None, display_name=None, description=None):
        """ Creates a channel within a specified team

        :param team_id: the team_id where the channel is created.

        :rtype: channel
        """

        if not team_id and display_name:
            raise RuntimeError('Provide the team_id and the display_name')

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_channels').format(team_id=team_id))

        if display_name and description:
            data = {
                'displayName': display_name,
                'description': description,
            }
        else:
            data = {
                'displayName': display_name,
            }

        response = self.con.post(url, data=data)

        if not response:
            return None

        data = response.json()

        return self.channel_constructor(parent=self, **{self._cloud_data_key: data})

    def get_channel_info(self, team_id=None, channel_id=None):
        """ Returns the channel info for a given channel

        :param team_id: the team_id of the channel to get the info of.
        :param channel_id: the channel_id of the channel to get the info of.

        :rtype: channel
        """

        if not team_id and channel_id:
            raise RuntimeError('Provide the team_id and channel_id')

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_channel_info').format(team_id=team_id, channel_id=channel_id))

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return self.channel_constructor(parent=self, **{self._cloud_data_key: data})

    def get_apps_in_team(self, team_id=None):
        """ Returns a list of apps of a specified team

        :param team_id: the team_id of the team to get the apps of.

        :rtype: apps
        """

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_apps_in_team').format(team_id=team_id))
        else:
            raise RuntimeError('Provide the team_id')

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return [
            self.app_constructor(
                parent=self, **{self._cloud_data_key: site})
            for site in data.get('value', [])]
<commit_msg>Fix team displayName -> display_name<commit_after>import logging

from dateutil.parser import parse
from .utils import ApiComponent

log = logging.getLogger(__name__)

class Team(ApiComponent):
    """ A Microsoft Teams team """

    _endpoints = {}

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Microsoft Teams team

        :param parent: parent object
        :type parent: Teams
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        cloud_data = kwargs.get(self._cloud_data_key, {})

        self.object_id = cloud_data.get('id')

        # Choose the main_resource passed in kwargs over parent main_resource
        main_resource = kwargs.pop('main_resource', None) or (
            getattr(parent, 'main_resource', None) if parent else None)

        main_resource = '{}{}'.format(main_resource, '')

        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

        self.display_name = cloud_data.get(self._cc('displayName'), '')
        self.description = cloud_data.get(self._cc('description'), '')
        self.is_archived = cloud_data.get(self._cc('isArchived'), '')
        self.web_url = cloud_data.get(self._cc('webUrl'), '')

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'Team: {}'.format(self.display_name)

    def __eq__(self, other):
        return self.object_id == other.object_id

class Channel(ApiComponent):
    """ A Microsoft Teams channel """

    _endpoints = {}

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Microsoft Teams channel

        :param parent: parent object
        :type parent: Teams
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        cloud_data = kwargs.get(self._cloud_data_key, {})

        self.object_id = cloud_data.get('id')

        # Choose the main_resource passed in kwargs over parent main_resource
        main_resource = kwargs.pop('main_resource', None) or (
            getattr(parent, 'main_resource', None) if parent else None)

        main_resource = '{}{}'.format(main_resource, '')

        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

        self.display_name = cloud_data.get(self._cc('displayName'), '')
        self.description = cloud_data.get('description')
        self.email = cloud_data.get('email')

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'Channel: {}'.format(self.display_name)

    def __eq__(self, other):
        return self.object_id == other.object_id

class App(ApiComponent):
    """ A Microsoft Teams app """

    _endpoints = {}

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Microsoft Teams app

        :param parent: parent object
        :type parent: Teams
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        cloud_data = kwargs.get(self._cloud_data_key, {})

        self.object_id = cloud_data.get('id')

        # Choose the main_resource passed in kwargs over parent main_resource
        main_resource = kwargs.pop('main_resource', None) or (
            getattr(parent, 'main_resource', None) if parent else None)

        main_resource = '{}{}'.format(main_resource, '')

        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

        self.app_definition = cloud_data.get(self._cc('teamsAppDefinition'), '')

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'App: {}'.format(self.app_definition['displayName'])

    def __eq__(self, other):
        return self.object_id == other.object_id

class Teams(ApiComponent):
    """ A microsoft teams class
        In order to use the API following permissions are required.
        Delegated (work or school account) - Group.Read.All, Group.ReadWrite.All
    """

    _endpoints = {
        'get_my_teams': '/me/joinedTeams',
        'get_channels': '/teams/{team_id}/channels',
        'create_channel': '/teams/{team_id}/channels',
        'get_channel_info': '/teams/{team_id}/channels/{channel_id}',
        'get_apps_in_team': '/teams/{team_id}/installedApps?$expand=teamsAppDefinition',
    }

    team_constructor = Team
    channel_constructor = Channel
    app_constructor = App

    def __init__(self, *, parent=None, con=None, **kwargs):
        """ A Teams object

        :param parent: parent object
        :type parent: Account
        :param Connection con: connection to use if no parent specified
        :param Protocol protocol: protocol to use if no parent specified
         (kwargs)
        :param str main_resource: use this resource instead of parent resource
         (kwargs)
        """
        if parent and con:
            raise ValueError('Need a parent or a connection but not both')
        self.con = parent.con if parent else con

        # Choose the main_resource passed in kwargs over the host_name
        main_resource = kwargs.pop('main_resource',
                                   '')  # defaults to blank resource
        super().__init__(
            protocol=parent.protocol if parent else kwargs.get('protocol'),
            main_resource=main_resource)

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return 'Microsoft Teams'

    def get_my_teams(self, *args):
        """ Returns a list of teams that I am in

        :rtype: teams
        """

        url = self.build_url(self._endpoints.get('get_my_teams'))

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return [
            self.team_constructor(parent=self, **{self._cloud_data_key: site})
            for site in data.get('value', [])]

    def get_channels(self, team_id=None):
        """ Returns a list of channels of a specified team

        :param team_id: the team_id of the channel to be retrieved.

        :rtype: channels
        """

        if not team_id:
            raise RuntimeError('Provide the team_id')

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_channels').format(team_id=team_id))

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return [
            self.channel_constructor(parent=self, **{self._cloud_data_key: site})
            for site in data.get('value', [])]

    def create_channel(self, team_id=None, display_name=None, description=None):
        """ Creates a channel within a specified team

        :param team_id: the team_id where the channel is created.

        :rtype: channel
        """

        if not team_id and display_name:
            raise RuntimeError('Provide the team_id and the display_name')

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_channels').format(team_id=team_id))

        if display_name and description:
            data = {
                'displayName': display_name,
                'description': description,
            }
        else:
            data = {
                'displayName': display_name,
            }

        response = self.con.post(url, data=data)

        if not response:
            return None

        data = response.json()

        return self.channel_constructor(parent=self, **{self._cloud_data_key: data})

    def get_channel_info(self, team_id=None, channel_id=None):
        """ Returns the channel info for a given channel

        :param team_id: the team_id of the channel to get the info of.
        :param channel_id: the channel_id of the channel to get the info of.

        :rtype: channel
        """

        if not team_id and channel_id:
            raise RuntimeError('Provide the team_id and channel_id')

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_channel_info').format(team_id=team_id, channel_id=channel_id))

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return self.channel_constructor(parent=self, **{self._cloud_data_key: data})

    def get_apps_in_team(self, team_id=None):
        """ Returns a list of apps of a specified team

        :param team_id: the team_id of the team to get the apps of.

        :rtype: apps
        """

        if team_id:
            # get channels by the team id
            url = self.build_url(
                self._endpoints.get('get_apps_in_team').format(team_id=team_id))
        else:
            raise RuntimeError('Provide the team_id')

        response = self.con.get(url)

        if not response:
            return None

        data = response.json()

        return [
            self.app_constructor(
                parent=self, **{self._cloud_data_key: site})
            for site in data.get('value', [])]
<|endoftext|>


<commit_before>"""This module handles lists"""

def de_duplicate(items):
    """Remove any duplicate item, preserving order

    >>> de_duplicate([1, 2, 1, 2])
    [1, 2]
    """
    result = []
    for item in items:
        if item not in result:
            result.append(item)
    return result
<commit_msg>Flatten lists<commit_after>"""This module handles lists"""

import itertools

def de_duplicate(items):
    """Remove any duplicate item, preserving order

    >>> assert de_duplicate([1, 9, 2, 8, 1, 7, 2]) == [1, 9, 2, 8, 7]
    """
    result = []
    for item in items:
        if item not in result:
            result.append(item)
    return result

def flatten(list_of_lists):
    """Reduce a lists of lists to a list, comprehensively

    >>> assert flatten([[1, 2], [3, 4]]) == [1, 2, 3, 4]

    Thanks to Alex Martelli at
        https://stackoverflow.com/a/952952/500942
    """
    return [item for list_ in list_of_lists for item in list_]

def flatten_(list_of_lists):
    """Reduce a lists of lists to a list, functionally

    >>> assert flatten_([[1, 2], [3, 4]]) == [1, 2, 3, 4]

    Thanks to CTT at
        https://stackoverflow.com/a/716482/500942
    """
    return list(itertools.chain.from_iterable(list_of_lists))

<|endoftext|>


<commit_before><commit_msg>calcs/hazard/classical/__init__:<commit_after><|endoftext|>


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

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# Copyright (c) 2014 Mozilla Corporation
#
# Contributors:
# Anthony Verez averez@mozilla.com
# Jeff Bryner jbryner@mozilla.com
# Michal Purzynski <mpurzynski@mozilla.com>

from lib.alerttask import AlertTask
from query_models import SearchQuery, TermMatch, ExistsMatch, TermsMatch

class AlertMultipleIntelHits(AlertTask):
    def main(self):
        search_query = SearchQuery(minutes=2)

        search_query.add_must([
            TermMatch('_type', 'bro'),
            TermMatch('eventsource', 'nsm'),
            TermMatch('category', 'brointel'),
            ExistsMatch('details.seenindicator'),
            TermsMatch('hostname', ['sensor1', 'sensor2', 'sensor3'])
        ])

        self.filtersManual(search_query)

        # Search aggregations on field 'seenindicator', keep X samples of events at most
        self.searchEventsAggregated('details.seenindicator', samplesLimit=10)
        # alert when >= X matching events in an aggregation
        self.walkAggregations(threshold=10)

    # Set alert properties
    def onAggregation(self, aggreg):
        # aggreg['count']: number of items in the aggregation, ex: number of failed login attempts
        # aggreg['value']: value of the aggregation field, ex: toto@example.com
        # aggreg['events']: list of events in the aggregation
        category = 'bro'
        tags = ['nsm,bro,intel']
        severity = 'NOTICE'
        hostname = aggreg['events'][0]['_source']['hostname']

        summary = '{0} {1} {2} on {3}'.format(aggreg['count'], hostname, ' Bro intel match for indicator:', aggreg['value'])

        summary += ' sample hosts that hit it: '
        for e in aggreg['events'][:3]:
            if 'details' in e['_source'].keys() \
               and 'sourceipaddress' in e['_source']['details'].keys() \
               and 'seenwhere' in e['_source']['details'].keys():
                interestingaddres = ''
                # someone talking to a bad guy, I want to know who
                # someone resolving bad guy's domain name, I want to know who
                # bad guy talking to someone, I want to know to whom
                if 'Conn::IN_RESP' in e['_source']['details']['seenwhere'] \
                    or 'HTTP::IN_HOST_HEADER' in e['_source']['details']['seenwhere'] \
                    or 'DNS::IN_REQUEST' in e['_source']['details']['seenwhere']:
                    interestingaddres = e['_source']['details']['sourceipaddress']
                elif 'Conn::IN_ORIG' in e['_source']['details']['seenwhere'] \
                    or 'HTTP::IN_X_CLUSTER_CLIENT_IP_HEADER' in e['_source']['details']['seenwhere'] \
                    or 'HTTP::IN_X_FORWARDED_FOR_HEADER' in e['_source']['details']['seenwhere']:
                    interestingaddres = e['_source']['details']['destinationipaddress']

                summary += '{0} in {1} '.format(interestingaddres, e['_source']['details']['seenwhere'])

        # Create the alert object based on these properties
        return self.createAlertDict(summary, category, tags, aggreg['events'], severity)
<commit_msg>Update the multiple-intel-iocs-detected alert<commit_after>#!/usr/bin/env python

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# Copyright (c) 2014 Mozilla Corporation
#
# Contributors:
# Anthony Verez averez@mozilla.com
# Jeff Bryner jbryner@mozilla.com
# Michal Purzynski <mpurzynski@mozilla.com>

from lib.alerttask import AlertTask
from query_models import SearchQuery, TermMatch, ExistsMatch, TermsMatch

class AlertMultipleIntelHits(AlertTask):
    def main(self):
        search_query = SearchQuery(minutes=2)

        search_query.add_must([
            TermMatch('_type', 'nsm'),
            TermMatch('category', 'bro'),
            TermMatch('source', 'intel'),
            ExistsMatch('details.seenindicator'),
            TermsMatch('hostname', ['sensor1', 'sensor2', 'sensor3'])
        ])

        self.filtersManual(search_query)

        # Search aggregations on field 'seenindicator', keep X samples of events at most
        self.searchEventsAggregated('details.seenindicator', samplesLimit=10)
        # alert when >= X matching events in an aggregation
        self.walkAggregations(threshold=10)

    # Set alert properties
    def onAggregation(self, aggreg):
        # aggreg['count']: number of items in the aggregation, ex: number of failed login attempts
        # aggreg['value']: value of the aggregation field, ex: toto@example.com
        # aggreg['events']: list of events in the aggregation
        category = 'bro'
        tags = ['nsm,bro,intel']
        severity = 'NOTICE'
        hostname = aggreg['events'][0]['_source']['hostname']

        summary = '{0} {1} {2} on {3}'.format(aggreg['count'], hostname, ' Bro intel match for indicator:', aggreg['value'])

        summary += ' sample hosts that hit it: '
        for e in aggreg['events'][:3]:
            if 'details' in e['_source'].keys() \
               and 'sourceipaddress' in e['_source']['details'].keys() \
               and 'seenwhere' in e['_source']['details'].keys():
                interestingaddres = ''
                # someone talking to a bad guy, I want to know who
                # someone resolving bad guy's domain name, I want to know who
                # bad guy talking to someone, I want to know to whom
                if 'Conn::IN_RESP' in e['_source']['details']['seenwhere'] \
                    or 'HTTP::IN_HOST_HEADER' in e['_source']['details']['seenwhere'] \
                    or 'DNS::IN_REQUEST' in e['_source']['details']['seenwhere']:
                    interestingaddres = e['_source']['details']['sourceipaddress']
                elif 'Conn::IN_ORIG' in e['_source']['details']['seenwhere'] \
                    or 'HTTP::IN_X_CLUSTER_CLIENT_IP_HEADER' in e['_source']['details']['seenwhere'] \
                    or 'HTTP::IN_X_FORWARDED_FOR_HEADER' in e['_source']['details']['seenwhere']:
                    interestingaddres = e['_source']['details']['destinationipaddress']

                summary += '{0} in {1} '.format(interestingaddres, e['_source']['details']['seenwhere'])

        # Create the alert object based on these properties
        return self.createAlertDict(summary, category, tags, aggreg['events'], severity)
<|endoftext|>


<commit_before>import znc

# ['akill', 'nick', 'time', 'reason', 'optional']
#    [0]      [1]    [2]       [3]        [4]

class akill(znc.Module):
    description = "Shortcut commands for akill"
    module_types = [znc.CModInfo.NetworkModule]
    reasons = ["netban", "spam", "given"]
    akill_format = "PRIVMSG OperServ :AKILL ADD +{duration} {mask} {reason}"
    echo_format = "AKILL ADD +{duration} {mask} {reason}"
    end_text = " Further violations of our network rules will result in an " \
               "increase in ban length and may become permanent. " \
               "https://snoonet.org/rules"

    def OnUserRaw(self, linecs):
        line = linecs.s.split()
        if line[0].lower() == "akill":
            if len(line) < 4:
                # too short, respond with error as a notice
                self.send_usage()
                return znc.HALT

            elif line[3][0] == "#" or line[3] in self.reasons:
                _, nick, time, reason, *address = line
                reason = reason.lower()
                address = " ".join(address)

                if line[3][0] == "#":
                    self.evasion(time, nick, reason, address or nick)
                elif reason == "netban":
                    self.netban(time, nick, address or nick)
                elif reason == "spam":
                    self.spam(time, nick, address or nick)
                elif reason == "given":
                    if address:
                        self.do_akill(nick, time, address)
                    else:
                        self.PutModNotice(
                            "given requires a reason to be specified")
                else:
                    self.send_usage()
            else:
                self.send_usage()
            return znc.HALT

    def evasion(self, time, nick, chan, address):
        reason = "{address}: {htime}network ban for ban evasion in " \
                 "{channel}.".format(
                    address=address, htime=self.human_time(time), channel=chan)
        self.do_akill(nick, time, reason + self.end_text)

    def netban(self, time, nick, address):
        reason = "{address}: {htime}network ban for evasion of a " \
                 "network ban.".format(
                    address=address, htime=self.human_time(time))
        self.do_akill(nick, time, reason + self.end_text)

    def spam(self, time, nick, address):
        reason = "{address}: {htime}network ban for spam. ".format(
                    address=address, htime=self.human_time(time))

        self.do_akill(nick, time, reason + self.end_text)

    def do_akill(self, nick, time, reason):
        self.PutModNotice(
            self.echo_format.format(duration=time, mask=nick, reason=reason)
        )

        self.PutIRC(
            self.akill_format.format(duration=time, mask=nick, reason=reason)
        )

    def send_usage(self):
        self.PutModNotice("Usage: /akill nick/mask time(defaults to days) "
                          "reason(Valid reasons are a channel name, spam, "
                          "netban and given. You specify the full message"
                          "with given) [nick to address]")

    @staticmethod
    def human_time(time: str):
        timechrs = {
            "m": " minute ",
            "h": " hour ",
            "d": " day ",
            "w": " week ",
            "y": " year "
        }
        formatted = False
        e_time = ""
        for char in time:
            if char in timechrs:
                e_time += timechrs[char]
                formatted = True
            else:
                e_time += char

        if not formatted:
            e_time += " day "
        return e_time
<commit_msg>Add the ability to add custom templates<commit_after>import json
from copy import copy

import znc

# ['akill', 'nick', 'time', 'reason', 'optional']
#    [0]      [1]    [2]       [3]        [4]

DEFAULT_REASONS = {
    "netban": "{address}: {htime} network ban for evasion of a network ban. {trail}",
    "spam": "{address}: {htime} network ban for spam. {trail}",
    "given": "{text}",
    "evasion": "{address}: {htime} network ban for ban evasion in {reason}. {trail}",
}

DEFAULT_FORMAT = "AKILL ADD +{duration} {mask} {reason}"
MSG_FORMAT = "PRIVMSG OperServ :{cmd}"

DEFAULT_TRAIL = "Further violations of our network rules will result in an " \
                "increase in ban length and may become permanent. " \
                "https://snoonet.org/rules"

def convert_table(data, headers):
    table = znc.CTable()
    for header in headers:
        table.AddColumn(header)

    for row in data:
        table.AddRow()
        for i, value in row:
            table.SetCell(headers[i], value)

    return table

class akill(znc.Module):
    description = "Shortcut commands for akill"
    module_types = [znc.CModInfo.NetworkModule]

    def OnLoad(self, args, msg):
        self.load_conf()
        return True

    def OnUserRaw(self, linecs):
        line = linecs.s.split()
        cmd = line[0].lower()
        if cmd != "akill":
            return znc.CONTINUE

        if len(line) < 4:
            # too short, respond with error as a notice
            self.send_usage()
            return znc.HALT

        _, nick, time, reason, *address = line
        reason = reason.lower()
        text = " ".join(address)
        address = ', '.join(address) or nick

        data = {
            "address": address,
            "htime": self.human_time(time),
            "nick": nick,
            "text": text,
            "reason": reason,
        }

        if reason[0] == "#":
            reason = "evasion"
        elif reason == "evasion":
            self.PutModNotice("You must specify a channel as the reason when banning for evasion.")

        if reason in self.reasons:
            fmt = self.reasons[reason]
            try:
                out = fmt.format_map(data)
            except LookupError:
                self.PutModNotice("Missing required parameter")
            else:
                self.do_akill(nick, time, out)
        else:
            self.send_usage()

        return znc.HALT

    def OnModCommand(self, text):
        cmd, _, args = str(text).partition(' ')
        cmd = cmd.lower()

        out = ""

        if cmd == "help":
            out = self.cmd_help()

        elif cmd == "listreasons":
            out = self.list_reasons()

        elif cmd == "setreason":
            out = self.add_reason(args)

        elif cmd == "delreason":
            out = self.del_reason(args)

        elif cmd == "settrail":
            out = self.set_trail(args)

        elif cmd == "setcmd":
            out = self.set_cmd(args)

        elif cmd == "setformat":
            out = self.set_format(args)

        else:
            out = "Unknown command. See Help for more information"

        if out:
            self.PutModule(out)

    def cmd_help(self):
        headers = ["Command", "Arguments", "Description"]
        table = (
            ("Help", "", "Displays the help for this module"),
            ("ListReasons", "", "Lists all currently configured reasons"),
            ("SetReason", "<name> <format>", "Sets the format for the specified reason"),
            ("DelReason", "<name>", "Removes a reason from the list"),
            ("SetTrail", "<trail>", "Sets the trailing text to be appended to akill reasons"),
            ("SetCmd", "<format>", "Set the format for the command to send (default: {})".format(DEFAULT_FORMAT)),
            ("SetFormat", "<format>",
             "Set the format for the raw line to send to the server (default: {})".format(MSG_FORMAT)),
        )

        return convert_table(table, headers)

    def list_reasons(self):
        headers = ["Name", "Format"]
        return convert_table(self.reasons.items(), headers)

    def add_reason(self, args):
        name, _, fmt = args.partition(' ')
        if not name:
            return "Missing required parameter: name"

        if not fmt:
            return "Missing required parameter: format"

        name = name.lower()
        old = self.reasons.get(name)
        self.reasons[name] = fmt
        self.save_conf()

        if old is not None:
            return "Format set. Old format was \"{}\"".format(old)

        return "Format set."

    def del_reason(self, args):
        name = args.strip().lower()

        if name not in self.reasons:
            return "Unknown reason."

        del self.reasons[name]
        self.save_conf()
        return "Reason deleted."

    def set_trail(self, args):
        self.trail = args
        self.save_conf()
        return "Trail set."

    def set_cmd(self, args):
        self.echo_format = args
        self.save_conf()
        return "Set command."

    def set_format(self, args):
        self.akill_format = args
        self.save_conf()
        return "Set raw format."

    def do_akill(self, nick, time, reason):
        cmd = self.echo_format.format(duration=time, mask=nick, reason=reason)
        self.PutModNotice(cmd)
        self.PutIRC(self.akill_format.format(cmd=cmd))

    def send_usage(self):
        self.PutModNotice("Usage: /akill <target> <duration (default unit: days)> <reason> [nicks to address]")

    def load_conf(self):
        reasons = self.nv.get("reasons")

        if reasons is None:
            self.reasons = copy(DEFAULT_REASONS)
        else:
            self.reasons = json.loads(reasons)

        self.echo_format = self.nv.get("echo_format", DEFAULT_FORMAT)
        self.akill_format = self.nv.get("akill_format", MSG_FORMAT)

        self.trail = self.nv.get("trail", DEFAULT_TRAIL)

    def save_conf(self):
        self.nv["reasons"] = json.dumps(self.reasons)
        self.nv["echo_format"] = self.echo_format
        self.nv["akill_format"] = self.akill_format
        self.nv["trail"] = self.trail

    @staticmethod
    def human_time(time: str):
        timechrs = {
            "m": "minute",
            "h": "hour",
            "d": "day",
            "w": "week",
            "y": "year",
        }
        parts = [""]
        for char in time:
            if char in timechrs:
                parts[-1] += " " + timechrs[char]
                parts.append("")
            else:
                parts[-1] += char

        if not parts[0]:
            parts[0] += " day"

        return ', '.join(parts)
<|endoftext|>