Python timeout_decorator.timeout() Examples

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() 

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) 

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) 

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 

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 

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 

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 

def not_test_dtls_scanner_sniff_pos(self):
        scanner = DTLSScanner(TestParams())
        scanner.sniff(timeout=0.0001) 

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 

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 

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 

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 

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)]) 

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 

def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 

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.") 

def ratio_timeout(self, ratio=0.5):
        return timeout_decorator.timeout(
            ratio * self.time_budget - (time.time() - self.start_time)
        ) 

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()) 

def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 

def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 

def __init__(self, error):
        if isinstance(error, dns.exception.Timeout):
            error.kwargs["timeout"] = round(error.kwargs["timeout"], 1) 

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 

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 

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 

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)]) 

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 

def checkFormat(binary_name,inputType="STDIN"):

    p = angr.Project(binary_name,load_options={"auto_load_libs": False})

    p.hook_symbol('printf',printFormat)

    #Setup state based on input type
    argv = [binary_name]
    if inputType == "STDIN":
        state = p.factory.full_init_state(args=argv)
    elif inputType == "LIBPWNABLE":
        handle_connection = p.loader.main_object.get_symbol('handle_connection')
        state = p.factory.entry_state(addr=handle_connection.rebased_addr)
    else:
        arg = claripy.BVS("arg1", 300 * 8)
        argv.append(arg)
        state = p.factory.full_init_state(args=argv)
        state.globals['arg'] = arg

    state.globals['inputType'] = inputType
    simgr = p.factory.simgr(state, immutable=False, save_unconstrained=True)

    run_environ = {}
    run_environ['type'] = None
    end_state = None
    #Lame way to do a timeout
    try:
        @timeout_decorator.timeout(120)
        def exploreBinary(simgr):
            simgr.explore(find=lambda s: 'type' in s.globals)

        exploreBinary(simgr)
        if 'found' in simgr.stashes and len(simgr.found):
            end_state = simgr.found[0]
            run_environ['type'] = end_state.globals['type']
            run_environ['position'] = end_state.globals['position']
            run_environ['length'] = end_state.globals['length']

    except (KeyboardInterrupt, timeout_decorator.TimeoutError) as e:
        print("[~] Format check timed out")
    if (inputType == "STDIN" or inputType == "LIBPWNABLE")and end_state is not None:
        stdin_str = str(end_state.posix.dumps(0))
        print("[+] Triggerable with STDIN : {}".format(stdin_str))
        run_environ['input'] = stdin_str
    elif inputType == "ARG" and end_state is not None:
        arg_str = str(end_state.solver.eval(arg,cast_to=str))
        run_environ['input'] = arg_str
        print("[+] Triggerable with arg : {}".format(arg_str))
       
    return run_environ 

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)]) 

def _PsDriveRemoteCopy(self, local_path, remote_path,
                         copy_to, network_drive):
    """Copies a file to or from the VM using New-PSDrive and Copy-Item.

    Args:
      local_path: Local path to file.
      remote_path: Optional path of where to copy file on remote host.
      copy_to: True to copy to vm, False to copy from vm.
      network_drive: The smb specification for the remote drive
          (//{ip_address}/{share_name}).

    Raises:
      RemoteCommandError: If there was a problem copying the file.
    """
    set_error_pref = '$ErrorActionPreference="Stop"'

    password = self.password.replace("'", "''")
    create_cred = (
        '$pw = convertto-securestring -AsPlainText -Force \'%s\';'
        '$cred = new-object -typename System.Management.Automation'
        '.PSCredential -argumentlist %s,$pw' % (password, self.user_name))

    psdrive_name = self.name
    create_psdrive = (
        'New-PSDrive -Name %s -PSProvider filesystem -Root '
        '%s -Credential $cred' % (psdrive_name, network_drive))

    remote_path = '%s:%s' % (psdrive_name, remote_path)
    if copy_to:
      from_path, to_path = local_path, remote_path
    else:
      from_path, to_path = remote_path, local_path

    copy_item = 'Copy-Item -Path %s -Destination %s' % (from_path, to_path)

    delete_connection = 'net use %s /delete' % network_drive

    cmd = ';'.join([set_error_pref, create_cred, create_psdrive,
                    copy_item, delete_connection])

    stdout, stderr, retcode = vm_util.IssueCommand(
        ['powershell', '-Command', cmd], timeout=None, raise_on_failure=False)

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

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)