1#!/usr/bin/env python
2# Copyright 2009 Google Inc. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Define and process configuration from command-line or config file."""
17
18__author__ = 'tstromberg@google.com (Thomas Stromberg)'
19
20import ConfigParser
21import optparse
22import os.path
23import StringIO
24import tempfile
25
26# from third_party
27import httplib2
28
29import data_sources
30import util
31import addr_util
32import version
33
34SANITY_REFERENCE_URL = 'http://namebench.googlecode.com/svn/trunk/config/hostname_reference.cfg'
35
36
37def GetConfiguration(filename='config/namebench.cfg'):
38  """Get all of our configuration setup, args and config file."""
39  (options, args) = DefineAndParseOptions(filename=filename)
40  (configured_options, global_ns, regional_ns) = ProcessConfigurationFile(options)
41  supplied_ns = addr_util.ExtractIPTuplesFromString(' '.join(args))
42  return (configured_options, supplied_ns, global_ns, regional_ns)
43
44
45def DefineAndParseOptions(filename):
46  """Get our option configuration setup.
47
48  Args:
49    filename: path to configuration (may be relative)
50
51  Returns:
52    stuple of (OptionParser object, args)
53  """
54  ds = data_sources.DataSources()
55  import_types = ds.ListSourceTypes()
56  parser = optparse.OptionParser()
57  parser.add_option('-r', '--runs', dest='run_count', default=1, type='int',
58                    help='Number of test runs to perform on each nameserver.')
59  parser.add_option('-z', '--config', dest='config', default=filename,
60                    help='Config file to use.')
61  parser.add_option('-o', '--output', dest='output_file', default=None,
62                    help='Filename to write output to')
63  parser.add_option('-t', '--template', dest='template', default='html',
64                    help='Template to use for output generation (ascii, html, resolv.conf)')
65  parser.add_option('-c', '--csv_output', dest='csv_file', default=None,
66                    help='Filename to write query details to (CSV)')
67  parser.add_option('-j', '--health_threads', dest='health_thread_count', type='int',
68                    help='# of health check threads to use')
69  parser.add_option('-J', '--benchmark_threads', dest='benchmark_thread_count', type='int',
70                    help='# of benchmark threads to use')
71  parser.add_option('-P', '--ping_timeout', dest='ping_timeout', type='float',
72                    help='# of seconds ping requests timeout in.')
73  parser.add_option('-y', '--timeout', dest='timeout', type='float',
74                    help='# of seconds general requests timeout in.')
75  parser.add_option('-Y', '--health_timeout', dest='health_timeout',
76                    type='float', help='health check timeout (in seconds)')
77  parser.add_option('-i', '--input', dest='input_source',
78                    help=('Import hostnames from an filename or application (%s)' %
79                          ', '.join(import_types)))
80  parser.add_option('-I', '--invalidate_cache', dest='invalidate_cache',
81                    action='store_true',
82                    help='Force health cache to be invalidated')
83  parser.add_option('-q', '--query_count', dest='query_count', type='int',
84                    help='Number of queries per run.')
85  parser.add_option('-m', '--select_mode', dest='select_mode',
86                    default='automatic',
87                    help='Selection algorithm to use (weighted, random, chunk)')
88  parser.add_option('-s', '--num_servers', dest='num_servers',
89                    type='int', help='Number of nameservers to include in test')
90  parser.add_option('-S', '--system_only', dest='system_only',
91                    action='store_true', help='Only test current system nameservers.')
92  parser.add_option('-w', '--open_webbrowser', dest='open_webbrowser',
93                    action='store_true', help='Opens the final report in your browser')
94  parser.add_option('-u', '--upload_results', dest='upload_results',
95                    action='store_true', help='Upload anonymized results to SITE_URL (False)')
96  parser.add_option('-U', '--site_url', dest='site_url',
97                    help='URL to upload results to (http://namebench.appspot.com/)')
98  parser.add_option('-H', '--hide_results', dest='hide_results', action='store_true',
99                    help='Upload results, but keep them hidden from indexes.')
100  parser.add_option('-x', '--no_gui', dest='no_gui',
101                    action='store_true', help='Disable GUI')
102  parser.add_option('-C', '--enable-censorship-checks', dest='enable_censorship_checks',
103                    action='store_true', help='Enable censorship checks')
104  parser.add_option('-6', '--ipv6_only', dest='ipv6_only',
105                    action='store_true', help='Only include IPv6 name servers')
106  # Silly Mac OS X adding -psn_0_xxxx
107  parser.add_option('-p', '--psn')
108  parser.add_option('-O', '--only', dest='only',
109                    action='store_true',
110                    help='Only test nameservers passed as arguments')
111  return parser.parse_args()
112
113
114def GetLatestSanityChecks():
115  """Get the latest copy of the sanity checks config."""
116  h = httplib2.Http(tempfile.gettempdir(), timeout=10)
117  http_version_usable = False
118  use_config = None
119  content = None
120  try:
121    unused_resp, content = h.request(SANITY_REFERENCE_URL, 'GET')
122  except:
123    print '* Unable to fetch latest reference: %s' % util.GetLastExceptionString()
124  http_config = ConfigParser.ConfigParser()
125
126  if content and '[base]' in content:
127    fp = StringIO.StringIO(content)
128    try:
129      http_config.readfp(fp)
130      http_version_usable = True
131    except:
132      pass
133
134  ref_file = util.FindDataFile('config/hostname_reference.cfg')
135  local_config = ConfigParser.ConfigParser()
136  local_config.read(ref_file)
137
138  if http_version_usable:
139    if int(http_config.get('base', 'version')) > int(local_config.get('base', 'version')):
140      print '- Using %s' % SANITY_REFERENCE_URL
141      use_config = http_config
142
143  if not use_config:
144    use_config = local_config
145
146  return (use_config.items('sanity'),
147          use_config.items('sanity-secondary'),
148          use_config.items('censorship'))
149
150
151def ProcessConfigurationFile(options):
152  """Process configuration file, merge configuration with OptionParser.
153
154  Args:
155    options: optparse.OptionParser() object
156
157  Returns:
158    options: optparse.OptionParser() object
159    global_ns: A list of global nameserver tuples.
160    regional_ns: A list of regional nameservers tuples.
161
162  Raises:
163    ValueError: If we are unable to find a usable configuration file.
164  """
165  config = ConfigParser.ConfigParser()
166  full_path = util.FindDataFile(options.config)
167  config.read(full_path)
168  if not config or not config.has_section('general'):
169    raise ValueError('Could not find usable configuration in %s (%s)' % (full_path, options.config))
170  general = dict(config.items('general'))
171
172  if options.only or options.system_only:
173    global_ns = []
174    regional_ns = []
175  else:
176    global_ns = config.items('global')
177    regional_ns = config.items('regional') + config.items('private')
178
179  # -U implies -u
180  if options.site_url:
181    options.upload_results = True
182
183  for option in general:
184    if not getattr(options, option, None):
185      if 'timeout' in option:
186        value = float(general[option])
187      elif 'count' in option or 'num' in option or 'hide' in option:
188        value = int(general[option])
189      else:
190        value = general[option]
191      setattr(options, option, value)
192
193  for key in ('input_file', 'output_file', 'csv_file', 'input_source'):
194    value = getattr(options, key, None)
195    if value:
196      setattr(options, key, os.path.expanduser(value))
197
198  options.version = version.VERSION
199
200  return (options, global_ns, regional_ns)
201