Python timeout_decorator.timeout() Examples

The following are 30 code examples of timeout_decorator.timeout(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module timeout_decorator , or try the search function .
Example #1
Source Project: cotopaxi   Author: Samsung   File: common_test_utils.py    License: GNU General Public License v2.0 6 votes vote down vote up
def poke_tcp_server(server_port):
    for i in range(4):
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(0.1)
            sock.connect(("127.0.0.1", server_port))
            try:
                sock.send("123".encode(encoding="ascii"))
                print ("data sent (P3)")
            except (AttributeError, UnicodeDecodeError):
                sock.send(bytes("123"))
                print ("data sent (P2)")
            finally:
                sock.close()
                print ("socket closed")
        except (socket.timeout, socket.error):
            time.sleep(0.1) 
Example #2
Source Project: violent-python3   Author: BLTSEC   File: bot_net.py    License: GNU General Public License v3.0 6 votes vote down vote up
def connect(self):
        # this reduces the timeout of SSHClient's connect()
        # for failed dns lookups
        try:
            self.host_check(self.host)
        except Exception as e:
            # todo logging e
            print('[-] CONNECTION FAILED', self.host)
            return None
        try:
            s = SSHClient()
            s.set_missing_host_key_policy(AutoAddPolicy())
            if self.password:
                s.connect(self.host, port=22, username=self.user,
                          password=self.password, timeout=1)
            else:
                s.connect(self.host, port=22, username=self.user,
                          timeout=1, pkey=self.key)
            return s
        except Exception as e:
            # todo logging e
            print('[-] CONNECTION FAILED:', self.host) 
Example #3
Source Project: PerfKitBenchmarker   Author: GoogleCloudPlatform   File: windows_virtual_machine.py    License: Apache License 2.0 6 votes vote down vote up
def WaitForProcessRunning(self, process, timeout):
    """Blocks until either the timeout passes or the process is running.

    Args:
      process: string name of the process.
      timeout: number of seconds to block while the process is not running.

    Raises:
      WaitTimeoutError: raised if the process does not run within "timeout"
                        seconds.
    """
    command = ('$count={timeout};'
               'while( (ps | select-string {process} | measure-object).Count '
               '-eq 0 -and $count -gt 0) {{sleep 1; $count=$count-1}}; '
               'if ($count -eq 0) {{echo "FAIL"}}').format(
                   timeout=timeout, process=process)
    stdout, _ = self.RemoteCommand(command)
    if 'FAIL' in stdout:
      raise WaitTimeoutError() 
Example #4
Source Project: mobly   Author: google   File: monsoon.py    License: Apache License 2.0 5 votes vote down vote up
def measure_power(self, hz, duration, tag, offset=30):
        """Measure power consumption of the attached device.

        Because it takes some time for the device to calm down after the usb
        connection is cut, an offset is set for each measurement. The default
        is 30s. The total time taken to measure will be (duration + offset).

        Args:
            hz: Number of samples to take per second.
            duration: Number of seconds to take samples for in each step.
            offset: The number of seconds of initial data to discard.
            tag: A string that's the name of the collected data group.

        Returns:
            A MonsoonData object with the measured power data.
        """
        num = duration * hz
        oset = offset * hz
        data = None
        self.usb("auto")
        time.sleep(1)
        with self.dut.handle_usb_disconnect():
            time.sleep(1)
            try:
                data = self.take_samples(hz, num, sample_offset=oset)
                if not data:
                    raise MonsoonError(
                        "No data was collected in measurement %s." % tag)
                data.tag = tag
                self.dut.log.info("Measurement summary: %s", repr(data))
                return data
            finally:
                self.mon.StopDataCollection()
                self.log.info("Finished taking samples, reconnecting to dut.")
                self.usb("on")
                self.dut.adb.wait_for_device(timeout=DEFAULT_TIMEOUT_USB_ON)
                # Wait for device to come back online.
                time.sleep(10)
                self.dut.log.info("Dut reconnected.") 
Example #5
Source Project: KDD-Cup-AutoML-5   Author: pfnet-research   File: timer.py    License: MIT License 5 votes vote down vote up
def ratio_timeout(self, ratio=0.5):
        return timeout_decorator.timeout(
            ratio * self.time_budget - (time.time() - self.start_time)
        ) 
Example #6
Source Project: apm-integration-testing   Author: elastic   File: es.py    License: Apache License 2.0 5 votes vote down vote up
def es():
    class Elasticsearch(object):
        def __init__(self, url):
            verify = default.from_env("PYTHONHTTPSVERIFY") == "1"
            self.es = elasticsearch.Elasticsearch(url,
                                                  verify_certs=verify,
                                                  connection_class=elasticsearch.RequestsHttpConnection)
            self.index = "apm-*"

        def clean(self):
            self.es.indices.delete(self.index)
            self.es.indices.refresh()

        def term_q(self, filters):
            clauses = []
            for field, value in filters:
                if isinstance(value, list):
                    clause = {"terms": {field: value}}
                else:
                    clause = {"term": {field: {"value": value}}}
                clauses.append(clause)
            return {"query": {"bool": {"must": clauses}}}

        @timeout_decorator.timeout(10)
        def count(self, q):
            ct = 0
            while ct == 0:
                time.sleep(3)
                s = self.es.count(index=self.index, body=q)
                ct = s['count']
            return ct

    return Elasticsearch(getElasticsearchURL()) 
Example #7
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 
Example #8
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 
Example #9
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 
Example #10
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 
Example #11
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def _get_mx_hosts(domain, nameservers=None, timeout=2.0):
    """
    Queries DNS for a list of Mail Exchange hosts

    Args:
        domain (str): A domain name
        nameservers (list): A list of nameservers to query
        (Cloudflare's by default)

    Returns:
        list: A list of ``OrderedDicts``; each containing a ``preference``
                        integer and a ``hostname``

    Raises:
        :exc:`checkdmarc.DNSException`

    """
    hosts = []
    try:
        logging.debug("Checking for MX records on {0}".format(domain))
        answers = _query_dns(domain, "MX", nameservers=nameservers,
                             timeout=timeout)
        for record in answers:
            record = record.split(" ")
            preference = int(record[0])
            hostname = record[1].rstrip(".").strip().lower()
            hosts.append(OrderedDict(
                [("preference", preference), ("hostname", hostname)]))
        hosts = sorted(hosts, key=lambda h: (h["preference"], h["hostname"]))
    except dns.resolver.NXDOMAIN:
        raise DNSException("The domain {0} does not exist".format(domain))
    except dns.resolver.NoAnswer:
        pass
    except Exception as error:
        raise DNSException(error)
    return hosts 
Example #12
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def _get_a_records(domain, nameservers=None, timeout=2.0):
    """
    Queries DNS for A and AAAA records

    Args:
        domain (str): A domain name
        nameservers (list): A list of nameservers to query
        (Cloudflare's by default)
        timeout(float): number of seconds to wait for an answer from DNS

    Returns:
        list: A sorted list of IPv4 and IPv6 addresses

    Raises:
        :exc:`checkdmarc.DNSException`

    """
    qtypes = ["A", "AAAA"]
    addresses = []
    for qt in qtypes:
        try:
            addresses += _query_dns(domain, qt, nameservers=nameservers,
                                    timeout=timeout)
        except dns.resolver.NXDOMAIN:
            raise DNSException("The domain {0} does not exist".format(domain))
        except dns.resolver.NoAnswer:
            # Sometimes a domain will only have A or AAAA records, but not both
            pass
        except Exception as error:
            raise DNSException(error)

    addresses = sorted(addresses)
    return addresses 
Example #13
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def _get_txt_records(domain, nameservers=None, timeout=2.0):
    """
    Queries DNS for TXT records

    Args:
        domain (str): A domain name
        nameservers (list): A list of nameservers to query
        (Cloudflare's by default)
        timeout(float): number of seconds to wait for an answer from DNS

    Returns:
        list: A list of TXT records

     Raises:
        :exc:`checkdmarc.DNSException`

    """
    try:
        records = _query_dns(domain, "TXT", nameservers=nameservers,
                             timeout=timeout)
    except dns.resolver.NXDOMAIN:
        raise DNSException("The domain {0} does not exist".format(domain))
    except dns.resolver.NoAnswer:
        raise DNSException(
            "The domain {0} does not have any TXT records".format(domain))
    except Exception as error:
        raise DNSException(error)

    return records 
Example #14
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def get_spf_record(domain, nameservers=None, timeout=2.0):
    """
    Retrieves and parses an SPF record

    Args:
        domain (str): A domain name
        nameservers (list): A list of nameservers to query
        (Cloudflare's by default)
        timeout(float): Number of seconds to wait for an answer from DNS

    Returns:
        OrderedDict: An SPF record parsed by result

    Raises:
        :exc:`checkdmarc.SPFRecordNotFound`
        :exc:`checkdmarc.SPFIncludeLoop`
        :exc:`checkdmarc.SPFRedirectLoop`
        :exc:`checkdmarc.SPFSyntaxError`
        :exc:`checkdmarc.SPFTooManyDNSLookups`

    """
    record = query_spf_record(domain, nameservers=nameservers, timeout=timeout)
    record = record["record"]
    parsed_record = parse_spf_record(record, domain, nameservers=nameservers,
                                     timeout=timeout)
    parsed_record["record"] = record

    return parsed_record 
Example #15
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def get_nameservers(domain, approved_nameservers=None,
                    nameservers=None, timeout=2.0):
    """
    Gets a list of nameservers for a given domain

    Args:
        domain (str): A domain name
        approved_nameservers (list): A list of approved nameserver substrings
        nameservers (list): A list of nameservers to qu+ery
        (Cloudflare's by default)
        timeout(float): number of seconds to wait for an record from DNS

    Returns:
        Dict: A dictionary with the following keys:
              - ``hostnames`` - A list of nameserver hostnames
              - ``warnings``  - A list of warnings
    """
    logging.debug("Getting NS records on {0}".format(domain))
    warnings = []

    ns_records = _get_nameservers(domain, nameservers=nameservers,
                                  timeout=timeout)

    if approved_nameservers:
        approved_nameservers = list(map(lambda h: h.lower(),
                                        approved_nameservers))
    for nameserver in ns_records:
        if approved_nameservers:
            approved = False
            for approved_nameserver in approved_nameservers:
                if approved_nameserver in nameserver.lower():
                    approved = True
                    break
            if not approved:
                warnings.append("Unapproved nameserver: {0}".format(
                    nameserver
                ))

    return OrderedDict([("hostnames", ns_records), ("warnings", warnings)]) 
Example #16
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 5 votes vote down vote up
def test_dnssec(domain, nameservers=None, timeout=2.0):
    """
    Check for DNSSEC on the given domain

    Args:
        domain (str): The domain to check
        nameservers (list): A list of nameservers to query
        timeout (float): Timeout in seconds

    Returns:
        bool: DNSSEC status
    """
    if nameservers is None:
        nameservers = ["1.1.1.1", "1.0.0.1",
                       "2606:4700:4700::1111", "2606:4700:4700::1001",
                       ]
    request = dns.message.make_query(get_base_domain(domain),
                                     dns.rdatatype.NS,
                                     want_dnssec=True)
    for nameserver in nameservers:
        try:
            response = dns.query.udp(request, nameserver, timeout=timeout)
            if response is not None:
                for record in response.answer:
                    if record.rdtype == dns.rdatatype.RRSIG:
                        if response.flags & dns.flags.AD:
                            return True
        except Exception as e:
            logging.debug("DNSSEC query error: {0}".format(e))

    return False 
Example #17
Source Project: ad-versarial   Author: ftramer   File: fridolin.py    License: MIT License 5 votes vote down vote up
def build_driver(use_adb=False, proxy=None, lang=None, timeout=MAX_TIMEOUT, fullscreen=True):
    opts = webdriver.ChromeOptions()

    #opts.add_argument("load-extension=../misc/jsinj/")

    if use_adb:
        opts.add_argument("load-extension=../misc/adblocker/unpacked/3.21.0_0/")

    if proxy:
        opts.add_argument("proxy-server={}".format(proxy))

    #opts.add_argument("start-maximized")
    if fullscreen:
        opts.add_argument("start-fullscreen")

    opts.add_argument("disable-infobars") # we remove the ugly yellow notification "Chrome is being controlled by automated test software"
    #opts.add_argument("disable-web-security") # We need this to disable SOP and access iframe content

    if lang:
        #print lang
        opts.add_argument("lang={}".format(lang))

    driver = webdriver.Chrome(executable_path="/usr/lib/chromium-browser/chromedriver", chrome_options=opts)

    if timeout:
        driver.set_page_load_timeout(timeout)

    return driver 
Example #18
Source Project: language   Author: google-research   File: official_evaluation.py    License: Apache License 2.0 5 votes vote down vote up
def try_executing_query(prediction, cursor, case_sensitive=True, verbose=False):
  """Attempts to execute a SQL query against a database given a cursor."""
  exception_str = None

  prediction_str = prediction[:]
  prediction_str = prediction_str.replace(';', '').strip()
  print('Current prediction:' + prediction_str)

  try:
    if not case_sensitive:
      new_prediction = ''
      last_quote = ''
      for char in prediction:
        new_prediction += char
        if char in {'"', '\''} and not last_quote:
          last_quote = char
        elif char == last_quote:
          last_quote = ''
          new_prediction += ' COLLATE NOCASE'
      prediction = new_prediction

      if verbose:
        print('Executing case-insensitive query:')
        print(new_prediction)
    pred_results = timeout_execute(cursor, prediction)
  except timeout_decorator.timeout_decorator.TimeoutError:
    print('!time out!')
    pred_results = []
    exception_str = 'timeout'
  except (sqlite3.Warning, sqlite3.Error, sqlite3.DatabaseError,
          sqlite3.IntegrityError, sqlite3.ProgrammingError,
          sqlite3.OperationalError, sqlite3.NotSupportedError) as e:
    exception_str = str(e).lower()
    pred_results = []

  return pred_results, exception_str 
Example #19
Source Project: datmo   Author: datmo   File: test_project.py    License: MIT License 5 votes vote down vote up
def test_dashboard(self):
        # test dashboard command
        self.project_command.parse(["dashboard"])

        @timeout_decorator.timeout(10, use_signals=False)
        def timed_run(timed_run_result):
            if self.project_command.execute():
                return timed_run_result

        # Failure case not initialized
        timed_run_result = False
        try:
            timed_run_result = timed_run(timed_run_result)
        except timeout_decorator.timeout_decorator.TimeoutError:
            timed_run_result = True

        assert not timed_run_result

        # Success case after initialization
        self.project_command.parse(
            ["init", "--name", "foobar", "--description", "test model"])

        @self.project_command.cli_helper.input("\n")
        def dummy(self):
            return self.project_command.execute()

        _ = dummy(self)

        self.project_command.parse(["dashboard"])
        timed_run_result = False
        try:
            timed_run_result = timed_run(timed_run_result)
        except timeout_decorator.timeout_decorator.TimeoutError:
            timed_run_result = True

        assert timed_run_result 
Example #20
Source Project: dcos-e2e   Author: dcos   File: _wait_for_dcos.py    License: Apache License 2.0 5 votes vote down vote up
def _test_utils_wait_for_dcos(
    session: Union[DcosApiSession, EnterpriseApiSession],
) -> None:
    """
    Wait for DC/OS using DC/OS Test Utils.

    DC/OS Test Utils raises its own timeout, a ``retrying.RetryError``.
    We want to ignore this error and use our own timeouts, so we wrap this in
    our own retried function.
    """
    session.wait_for_dcos()  # type: ignore 
Example #21
Source Project: cotopaxi   Author: Samsung   File: test_active_scanner.py    License: GNU General Public License v2.0 5 votes vote down vote up
def not_test_dtls_scanner_sniff_pos(self):
        scanner = DTLSScanner(TestParams())
        scanner.sniff(timeout=0.0001) 
Example #22
Source Project: cotopaxi   Author: Samsung   File: test_vulnerability_tester.py    License: GNU General Public License v2.0 5 votes vote down vote up
def test_vuln_tester_information_neg(self):
        self.vulnerability_tester_negative("GstRtspServer", "2019-12-01", "TP-LINK_000")

    # @timeout_decorator.timeout(5)
    #  def test_vuln_tester_madmaze(self):

    # self.vulnerability_tester("madmaze_htcpcp", "2011-08-03", "MADMAZE-HTCPCP_000")
    # pass 
Example #23
Source Project: PerfKitBenchmarker   Author: GoogleCloudPlatform   File: windows_virtual_machine.py    License: Apache License 2.0 5 votes vote down vote up
def RemoteCommand(self, command, should_log=False, ignore_failure=False,
                    suppress_warning=False, timeout=None):
    """Runs a powershell command on the VM.

    Args:
      command: A valid powershell command.
      should_log: A boolean indicating whether the command result should be
          logged at the info level. Even if it is false, the results will
          still be logged at the debug level.
      ignore_failure: Ignore any failure if set to true.
      suppress_warning: Suppress the result logging from IssueCommand when the
          return code is non-zero.
      timeout: Float. A timeout in seconds for the command. If None is passed,
          no timeout is applied. Timeout kills the winrm session which then
          kills the process being executed.

    Returns:
      A tuple of stdout and stderr from running the command.

    Raises:
      RemoteCommandError: If there was a problem issuing the command or the
          command timed out.
    """
    logging.info('Running command on %s: %s', self, command)
    s = winrm.Session(
        'https://%s:%s' % (self.GetConnectionIp(), self.winrm_port),
        auth=(self.user_name, self.password),
        server_cert_validation='ignore')
    encoded_command = six.ensure_str(
        base64.b64encode(command.encode('utf_16_le')))

    @timeout_decorator.timeout(timeout, use_signals=False,
                               timeout_exception=errors.VirtualMachine.
                               RemoteCommandError)
    def run_command():
      return s.run_cmd('powershell -encodedcommand %s' % encoded_command)

    r = run_command()
    retcode, stdout, stderr = r.status_code, six.ensure_str(
        r.std_out), six.ensure_str(r.std_err)

    debug_text = ('Ran %s on %s. Return code (%s).\nSTDOUT: %s\nSTDERR: %s' %
                  (command, self, retcode, stdout, stderr))
    if should_log or (retcode and not suppress_warning):
      logging.info(debug_text)
    else:
      logging.debug(debug_text)

    if retcode and not ignore_failure:
      error_text = ('Got non-zero return code (%s) executing %s\n'
                    'STDOUT: %sSTDERR: %s' %
                    (retcode, command, stdout, stderr))
      raise errors.VirtualMachine.RemoteCommandError(error_text)

    return stdout, stderr 
Example #24
Source Project: Zeratool   Author: ChrisTheCoolHut   File: formatExploiter.py    License: GNU General Public License v3.0 5 votes vote down vote up
def runIteration(binary_name,str_input,remote_server=False,remote_url="",remote_port=0,input_type="STDIN"):
    
    if input_type == "STDIN" or input_type == "LIBPWNABLE":
        if remote_server:
            proc = remote(remote_url,remote_port)
        else:
            proc = process(binary_name)
        proc.sendline(str_input)

        results = proc.recvall(timeout=5)
        print(results)
        results_split = results.split('_')

        #Get only hex strings of 8 characters or fewer
        position_leak = [x for x in results_split if len(x) < 9 and all([y in string.hexdigits for y in x])]

        #There should only be one
        leak = [position_leak][0]
    else:
        proc = process([binary_name,str_input.rstrip('\x00')])
    
        results = proc.recvall(timeout=5)
        print(results)
        results_split = results.split('_')

        #Get only hex strings of 8 characters or fewer
        position_leak = [x for x in results_split if len(x) < 9 and all([y in string.hexdigits for y in x])]

        #There should only be one
        leak = [position_leak][0]

    
    return leak 
Example #25
Source Project: mobly   Author: google   File: monsoon.py    License: Apache License 2.0 4 votes vote down vote up
def CollectData(self):
        """Return some current samples. Call StartDataCollection() first.
        """
        while 1:  # loop until we get data or a timeout
            _bytes = self._ReadPacket()
            if not _bytes:
                return None
            if len(_bytes) < 4 + 8 + 1 or _bytes[0] < 0x20 or _bytes[0] > 0x2F:
                logging.warning("Wanted data, dropped type=0x%02x, len=%d",
                                _bytes[0], len(_bytes))
                continue

            seq, _type, x, y = struct.unpack("BBBB", _bytes[:4])
            data = [
                struct.unpack(">hhhh", _bytes[x:x + 8])
                for x in range(4,
                               len(_bytes) - 8, 8)
            ]

            if self._last_seq and seq & 0xF != (self._last_seq + 1) & 0xF:
                logging.warning("Data sequence skipped, lost packet?")
            self._last_seq = seq

            if _type == 0:
                if not self._coarse_scale or not self._fine_scale:
                    logging.warning(
                        "Waiting for calibration, dropped data packet.")
                    continue
                out = []
                for main, usb, aux, voltage in data:
                    if main & 1:
                        coarse = ((main & ~1) - self._coarse_zero)
                        out.append(coarse * self._coarse_scale)
                    else:
                        out.append((main - self._fine_zero) * self._fine_scale)
                return out
            elif _type == 1:
                self._fine_zero = data[0][0]
                self._coarse_zero = data[1][0]
            elif _type == 2:
                self._fine_ref = data[0][0]
                self._coarse_ref = data[1][0]
            else:
                logging.warning("Discarding data packet type=0x%02x", _type)
                continue

            # See http://wiki/Main/MonsoonProtocol for details on these values.
            if self._coarse_ref != self._coarse_zero:
                self._coarse_scale = 2.88 / (self._coarse_ref -
                                             self._coarse_zero)
            if self._fine_ref != self._fine_zero:
                self._fine_scale = 0.0332 / (self._fine_ref - self._fine_zero) 
Example #26
Source Project: apm-integration-testing   Author: elastic   File: concurrent_requests.py    License: Apache License 2.0 4 votes vote down vote up
def check_counts(self, it, max_wait=60, backoff=1):
        err = "queried for {}, expected {}, got {}"

        def assert_count(terms, expected):
            """wait a bit for doc count to reach expectation"""
            @timeout_decorator.timeout(max_wait)
            def check_count(mut_actual):
                while True:
                    rsp = self.es.count(index=self.index, body=self.elasticsearch.term_q(terms))
                    mut_actual[0] = rsp["count"]
                    if mut_actual[0] >= expected:
                        return
                    time.sleep(backoff)

            mut_actual = [-1]  # keep actual count in this mutable
            try:
                check_count(mut_actual)
            except timeout_decorator.TimeoutError:
                pass
            actual = mut_actual[0]
            assert actual == expected, err.format(terms, expected, actual)

        self.es.indices.refresh()

        service_names = [ep.app_name for ep in self.endpoints]
        transactions_count = self.count("transaction") * it
        assert_count([("processor.event", "transaction"), ("service.name", service_names)], transactions_count)

        spans_count = self.count("span") * it
        assert_count([("processor.event", "span"), ('service.name', service_names)], spans_count)

        transactions_sum = spans_sum = 0
        for ep in self.endpoints:
            for span_name in ep.span_names:
                count = ep.count("span") * it / len(ep.span_names)
                spans_sum += count
                assert_count([
                    ("processor.event", "span"),
                    ("span.name", span_name),
                    ("service.name", ep.app_name),
                ], count)

            count = ep.count("transaction") * it
            transactions_sum += count
            assert_count([
                ("processor.event", "transaction"),
                ("service.name", ep.app_name),
                ("transaction.name", ep.transaction_name),
            ], count)

        assert transactions_count == transactions_sum, err.format(
            "transactions all endpoints", transactions_count, transactions_sum)
        assert spans_count == spans_sum, err.format(
            "spans all endpoints", spans_count, spans_sum) 
Example #27
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 4 votes vote down vote up
def _query_dns(domain, record_type, nameservers=None, timeout=2.0,
               cache=None):
    """
    Queries DNS

    Args:
        domain (str): The domain or subdomain to query about
        record_type (str): The record type to query for
        nameservers (list): A list of one or more nameservers to use
        (Cloudflare's public DNS resolvers by default)
        timeout (float): Sets the DNS timeout in seconds
        cache (ExpiringDict): Cache storage

    Returns:
        list: A list of answers
    """
    domain = str(domain).lower()
    record_type = record_type.upper()
    cache_key = "{0}_{1}".format(domain, record_type)
    if cache is None:
        cache = DNS_CACHE
    if cache:
        records = cache.get(cache_key, None)
        if records:
            return records

    resolver = dns.resolver.Resolver()
    timeout = float(timeout)
    if nameservers is None:
        nameservers = ["1.1.1.1", "1.0.0.1",
                       "2606:4700:4700::1111", "2606:4700:4700::1001",
                       ]
    resolver.nameservers = nameservers
    resolver.timeout = timeout
    resolver.lifetime = timeout
    if record_type == "TXT":
        resource_records = list(map(
            lambda r: r.strings,
            resolver.query(domain, record_type, lifetime=timeout)))
        _resource_record = [
            resource_record[0][:0].join(resource_record)
            for resource_record in resource_records if resource_record]
        records = [r.decode() for r in _resource_record]
    else:
        records = list(map(
            lambda r: r.to_text().replace('"', '').rstrip("."),
            resolver.query(domain, record_type, lifetime=timeout)))
    if cache:
        cache[cache_key] = records

    return records 
Example #28
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 4 votes vote down vote up
def query_dmarc_record(domain, nameservers=None, timeout=2.0):
    """
    Queries DNS for a DMARC record

    Args:
        domain (str): A domain name
        nameservers (list): A list of nameservers to query
        (Cloudflare's by default)
        timeout(float): number of seconds to wait for an record from DNS

    Returns:
        OrderedDict: An ``OrderedDict`` with the following keys:
                     - ``record`` - the unparsed DMARC record string
                     - ``location`` - the domain where the record was found
                     - ``warnings`` - warning conditions found

     Raises:
        :exc:`checkdmarc.DMARCRecordNotFound`
        :exc:`checkdmarc.DMARCRecordInWrongLocation`
        :exc:`checkdmarc.MultipleDMARCRecords`
        :exc:`checkdmarc.SPFRecordFoundWhereDMARCRecordShouldBe`

    """
    logging.debug("Checking for a DMARC record on {0}".format(domain))
    warnings = []
    base_domain = get_base_domain(domain)
    location = domain.lower()
    record = _query_dmarc_record(domain, nameservers=nameservers,
                                 timeout=timeout)
    try:
        root_records = _query_dns(domain.lower(), "TXT",
                                  nameservers=nameservers,
                                  timeout=timeout)
        for root_record in root_records:
            if root_record.startswith("v=DMARC1"):
                warnings.append("DMARC record at root of {0} "
                                "has no effect".format(domain.lower()))
    except Exception:
        pass

    if record is None and domain != base_domain:
        record = _query_dmarc_record(base_domain, nameservers=nameservers,
                                     timeout=timeout)
        location = base_domain
    if record is None:
        raise DMARCRecordNotFound(
            "A DMARC record does not exist for this domain or its base domain")

    return OrderedDict([("record", record), ("location", location),
                        ("warnings", warnings)]) 
Example #29
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 4 votes vote down vote up
def query_bimi_record(domain, selector="default", nameservers=None,
                      timeout=2.0):
    """
    Queries DNS for a BIMI record

    Args:
        domain (str): A domain name
        selector (str): The BMI selector
        nameservers (list): A list of nameservers to query
        (Cloudflare's by default)
        timeout(float): number of seconds to wait for an record from DNS

    Returns:
        OrderedDict: An ``OrderedDict`` with the following keys:
                     - ``record`` - the unparsed DMARC record string
                     - ``location`` - the domain where the record was found
                     - ``warnings`` - warning conditions found

     Raises:
        :exc:`checkdmarc.BIMIRecordNotFound`
        :exc:`checkdmarc.BIMIRecordInWrongLocation`
        :exc:`checkdmarc.MultipleBIMIRecords`

    """
    logging.debug("Checking for a BIMI record on {0}".format(domain))
    warnings = []
    base_domain = get_base_domain(domain)
    location = domain.lower()
    record = _query_bmi_record(domain, selector=selector,
                               nameservers=nameservers, timeout=timeout)
    try:
        root_records = _query_dns(domain.lower(), "TXT",
                                  nameservers=nameservers,
                                  timeout=timeout)
        for root_record in root_records:
            if root_record.startswith("v=BIMI1"):
                warnings.append("BIMI record at root of {0} "
                                "has no effect".format(domain.lower()))
    except Exception:
        pass

    if record is None and domain != base_domain and selector != "default":
        record = _query_bmi_record(base_domain, selector="default",
                                   nameservers=nameservers,
                                   timeout=timeout)
        location = base_domain
    if record is None:
        raise BIMIRecordNotFound(
            "A BIMI record does not exist for this domain or its base domain")

    return OrderedDict([("record", record), ("location", location),
                        ("warnings", warnings)]) 
Example #30
Source Project: checkdmarc   Author: domainaware   File: checkdmarc.py    License: Apache License 2.0 4 votes vote down vote up
def get_dmarc_record(domain, include_tag_descriptions=False, nameservers=None,
                     timeout=2.0):
    """
    Retrieves a DMARC record for a domain and parses it

    Args:
        domain (str): A domain name
        include_tag_descriptions (bool): Include descriptions in parsed results
        nameservers (list): A list of nameservers to query
        (Cloudflare's by default)
        timeout(float): number of seconds to wait for an answer from DNS

    Returns:
        OrderedDict: An ``OrderedDict`` with the following keys:
         - ``record`` - The DMARC record string
         - ``location`` -  Where the DMARC was found
         - ``parsed`` - See :meth:`checkdmarc.parse_dmarc_record`

     Raises:
        :exc:`checkdmarc.DMARCRecordNotFound`
        :exc:`checkdmarc.DMARCRecordInWrongLocation`
        :exc:`checkdmarc.MultipleDMARCRecords`
        :exc:`checkdmarc.SPFRecordFoundWhereDMARCRecordShouldBe`
        :exc:`checkdmarc.UnverifiedDMARCURIDestination`
        :exc:`checkdmarc.DMARCSyntaxError`
        :exc:`checkdmarc.InvalidDMARCTag`
        :exc:`checkdmarc.InvalidDMARCTagValue`
        :exc:`checkdmarc.InvalidDMARCReportURI`
        :exc:`checkdmarc.UnverifiedDMARCURIDestination`
        :exc:`checkdmarc.UnrelatedTXTRecordFound`
        :exc:`checkdmarc.DMARCReportEmailAddressMissingMXRecords`
    """
    query = query_dmarc_record(domain, nameservers=nameservers,
                               timeout=timeout)

    tag_descriptions = include_tag_descriptions

    tags = parse_dmarc_record(query["record"], query["location"],
                              include_tag_descriptions=tag_descriptions,
                              nameservers=nameservers, timeout=timeout)

    return OrderedDict([("record",
                         query["record"]),
                        ("location", query["location"]),
                        ("parsed", tags)])