script/traffic_replay

   1 #!/usr/bin/env python
   2 # Generates samba network traffic
   3 #
   4 # Copyright (C) Catalyst IT Ltd. 2017
   5 #
   6 # This program is free software; you can redistribute it and/or modify
   7 # it under the terms of the GNU General Public License as published by
   8 # the Free Software Foundation; either version 3 of the License, or
   9 # (at your option) any later version.
  10 #
  11 # This program is distributed in the hope that it will be useful,
  12 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14 # GNU General Public License for more details.
  15 #
  16 # You should have received a copy of the GNU General Public License
  17 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  18 #
  19 import sys
  20 import os
  21 import optparse
  22 import tempfile
  23 import shutil
  24
  25 sys.path.insert(0, "bin/python")
  26
  27 from samba.emulate import traffic
  28 import samba.getopt as options
  29
  30
  31 def main():
  32
  33     desc = ("Generates network traffic 'conversations' based on <summary-file>"
  34             " (which should the output file produced by either traffic_learner"
  35             " or traffic_summary.pl). This traffic is sent to <dns-hostname>,"
  36             " which is the full DNS hostname of the DC being tested.")
  37
  38     parser = optparse.OptionParser(
  39         "%prog [--help|options] <summary-file> <dns-hostname>",
  40         description=desc)
  41
  42     parser.add_option('--dns-rate', type='float', default=0,
  43                       help='fire extra DNS packets at this rate')
  44     parser.add_option('-B', '--badpassword-frequency',
  45                       type='float', default=0.0,
  46                       help='frequency of connections with bad passwords')
  47     parser.add_option('-K', '--prefer-kerberos',
  48                       action="store_true",
  49                       help='prefer kerberos when authenticating test users')
  50     parser.add_option('-I', '--instance-id', type='int', default=0,
  51                       help='Instance number, when running multiple instances')
  52     parser.add_option('-t', '--timing-data',
  53                       help=('write individual message timing data here '
  54                             '(- for stdout)'))
  55     parser.add_option('--preserve-tempdir', default=False, action="store_true",
  56                       help='do not delete temporary files')
  57     parser.add_option('-F', '--fixed-password',
  58                       type='string', default=None,
  59                       help=('Password used for the test users created. '
  60                             'Required'))
  61     parser.add_option('-c', '--clean-up',
  62                       action="store_true",
  63                       help='Clean up the generated groups and user accounts')
  64
  65     model_group = optparse.OptionGroup(parser, 'Traffic Model Options',
  66                                        'These options alter the traffic '
  67                                        'generated when the summary-file is a '
  68                                        'traffic-model (produced by '
  69                                        'traffic_learner)')
  70     model_group.add_option('-S', '--scale-traffic', type='float', default=1.0,
  71                            help='Increase the number of conversations by '
  72                            'this factor')
  73     model_group.add_option('-D', '--duration', type='float', default=None,
  74                            help=('Run model for this long (approx). '
  75                                  'Default 60s for models'))
  76     model_group.add_option('-r', '--replay-rate', type='float', default=1.0,
  77                            help='Replay the traffic faster by this factor')
  78     model_group.add_option('--traffic-summary',
  79                            help=('Generate a traffic summary file and write '
  80                                  'it here (- for stdout)'))
  81     parser.add_option_group(model_group)
  82
  83     user_gen_group = optparse.OptionGroup(parser, 'Generate User Options',
  84                                           "Add extra user/groups on the DC to "
  85                                           "increase the DB size. These extra "
  86                                           "users aren't used for traffic "
  87                                           "generation.")
  88     user_gen_group.add_option('-G', '--generate-users-only',
  89                               action="store_true",
  90                               help='Generate the users, but do not replay '
  91                               'the traffic')
  92     user_gen_group.add_option('-n', '--number-of-users', type='int', default=0,
  93                               help='Total number of test users to create')
  94     user_gen_group.add_option('--number-of-groups', type='int', default=0,
  95                               help='Create this many groups')
  96     user_gen_group.add_option('--average-groups-per-user',
  97                               type='int', default=0,
  98                               help='Assign the test users to this '
  99                               'many groups on average')
 100     user_gen_group.add_option('--group-memberships', type='int', default=0,
 101                               help='Total memberships to assign across all '
 102                               'test users and all groups')
 103     parser.add_option_group(user_gen_group)
 104
 105     sambaopts = options.SambaOptions(parser)
 106     parser.add_option_group(sambaopts)
 107     parser.add_option_group(options.VersionOptions(parser))
 108     credopts = options.CredentialsOptions(parser)
 109     parser.add_option_group(credopts)
 110
 111     # the --no-password credential doesn't make sense for this tool
 112     if parser.has_option('-N'):
 113         parser.remove_option('-N')
 114
 115     opts, args = parser.parse_args()
 116
 117     # First ensure we have reasonable arguments
 118
 119     if len(args) == 1:
 120         summary = None
 121         host    = args[0]
 122     elif len(args) == 2:
 123         summary, host = args
 124     else:
 125         parser.print_usage()
 126         return
 127
 128     if opts.clean_up:
 129         print >>sys.stderr, "Removing user and machine accounts"
 130         lp    = sambaopts.get_loadparm()
 131         creds = credopts.get_credentials(lp)
 132         ldb   = traffic.openLdb(host, creds, lp)
 133         traffic.clean_up_accounts(ldb, opts.instance_id)
 134         exit(0)
 135
 136     if summary:
 137         if not os.path.exists(summary):
 138             print >>sys.stderr, "Summary file %s doesn't exist" % summary
 139             sys.exit(1)
 140     # the summary-file can be ommitted for --generate-users-only and
 141     # --cleanup-up, but it should be specified in all other cases
 142     elif not opts.generate_users_only:
 143         print >>sys.stderr, "No summary-file specified to replay traffic from"
 144         sys.exit(1)
 145
 146     if not opts.fixed_password:
 147         print >>sys.stderr, ("Please use --fixed-password to specify a password"
 148                              " for the users created as part of this test")
 149         sys.exit(1)
 150
 151     lp = sambaopts.get_loadparm()
 152     creds = credopts.get_credentials(lp)
 153
 154     domain = opts.workgroup
 155     if domain:
 156         lp.set("workgroup", domain)
 157     else:
 158         domain = lp.get("workgroup")
 159         if domain == "WORKGROUP":
 160             print >>sys.stderr, ("NETBIOS domain does not appear to be "
 161                                  "specified, use the --workgroup option")
 162             sys.exit(1)
 163
 164     if not opts.realm and not lp.get('realm'):
 165         print >>sys.stderr, "Realm not specified, use the --realm option"
 166         sys.exit(1)
 167
 168     if opts.generate_users_only and not (opts.number_of_users or
 169                                          opts.number_of_groups):
 170         print >>sys.stderr, ("Please specify the number of users and/or groups "
 171                              "to generate.")
 172         sys.exit(1)
 173
 174     if opts.group_memberships and opts.average_groups_per_user:
 175         print >>sys.stderr, ("--group-memberships and --average-groups-per-user"
 176                              " are incompatible options - use one or the other")
 177         sys.exit(1)
 178
 179     if not opts.number_of_groups and opts.average_groups_per_user:
 180         print >>sys.stderr, ("--average-groups-per-user requires "
 181                              "--number-of-groups")
 182         sys.exit(1)
 183
 184     if not opts.number_of_groups and opts.group_memberships:
 185         print >>sys.stderr, "--group-memberships requires --number-of-groups"
 186         sys.exit(1)
 187
 188     if opts.timing_data not in ('-', None):
 189         try:
 190             open(opts.timing_data, 'w').close()
 191         except IOError as e:
 192             print >> sys.stderr, ("the supplied timing data destination "
 193                                   "(%s) is not writable" % opts.timing_data)
 194             print >> sys.stderr, e
 195             sys.exit()
 196
 197     if opts.traffic_summary not in ('-', None):
 198         try:
 199             open(opts.traffic_summary, 'w').close()
 200         except IOError as e:
 201             print >> sys.stderr, ("the supplied traffic summary destination "
 202                                   "(%s) is not writable" % opts.traffic_summary)
 203             print >> sys.stderr, e
 204             sys.exit()
 205
 206     traffic.DEBUG_LEVEL = opts.debuglevel
 207
 208     duration = opts.duration
 209     if duration is None:
 210         duration = 60.0
 211
 212     # ingest the model or traffic summary
 213     if summary:
 214         try:
 215             conversations, interval, duration, dns_counts = \
 216                                             traffic.ingest_summaries([summary])
 217
 218             print >>sys.stderr, ("Using conversations from the traffic summary "
 219                                  "file specified")
 220
 221             # honour the specified duration if it's different to the
 222             # capture duration
 223             if opts.duration is not None:
 224                 duration = opts.duration
 225
 226         except ValueError as e:
 227             if not e.message.startswith('need more than'):
 228                 raise
 229
 230             model = traffic.TrafficModel()
 231
 232             try:
 233                 model.load(summary)
 234             except ValueError:
 235                 print >>sys.stderr, ("Could not parse %s. The summary file "
 236                                      "should be the output from either the "
 237                                      "traffic_summary.pl or "
 238                                      "traffic_learner scripts."
 239                                      % summary)
 240                 sys.exit()
 241
 242             print >>sys.stderr, ("Using the specified model file to "
 243                                  "generate conversations")
 244
 245             conversations = model.generate_conversations(opts.scale_traffic,
 246                                                          duration,
 247                                                          opts.replay_rate)
 248
 249     else:
 250         conversations = []
 251
 252     if opts.debuglevel > 5:
 253         for c in conversations:
 254             for p in c.packets:
 255                 print "    ", p
 256
 257         print '=' * 72
 258
 259     if opts.number_of_users and opts.number_of_users < len(conversations):
 260         print >>sys.stderr, ("--number-of-users (%d) is less than the "
 261                              "number of conversations to replay (%d)"
 262                              % (opts.number_of_users, len(conversations)))
 263         sys.exit(1)
 264
 265     number_of_users = max(opts.number_of_users, len(conversations))
 266     max_memberships = number_of_users * opts.number_of_groups
 267
 268     if not opts.group_memberships and opts.average_groups_per_user:
 269         opts.group_memberships = opts.average_groups_per_user * number_of_users
 270         print >>sys.stderr, ("Using %d group-memberships based on %u average "
 271                              "memberships for %d users"
 272                              % (opts.group_memberships,
 273                                 opts.average_groups_per_user, number_of_users))
 274
 275     if opts.group_memberships > max_memberships:
 276         print >>sys.stderr, ("The group memberships specified (%d) exceeds "
 277                              "the total users (%d) * total groups (%d)"
 278                              % (opts.group_memberships, number_of_users,
 279                                 opts.number_of_groups))
 280         sys.exit(1)
 281
 282     try:
 283         ldb = traffic.openLdb(host, creds, lp)
 284     except:
 285         print >>sys.stderr, ("\nInitial LDAP connection failed! Did you supply "
 286                              "a DNS host name and the correct credentials?")
 287         sys.exit(1)
 288
 289     if opts.generate_users_only:
 290         traffic.generate_users_and_groups(ldb,
 291                                           opts.instance_id,
 292                                           opts.fixed_password,
 293                                           opts.number_of_users,
 294                                           opts.number_of_groups,
 295                                           opts.group_memberships)
 296         sys.exit()
 297
 298     tempdir = tempfile.mkdtemp(prefix="samba_tg_")
 299     print >>sys.stderr, "Using temp dir %s" % tempdir
 300
 301     traffic.generate_users_and_groups(ldb,
 302                                       opts.instance_id,
 303                                       opts.fixed_password,
 304                                       number_of_users,
 305                                       opts.number_of_groups,
 306                                       opts.group_memberships)
 307
 308     accounts = traffic.generate_replay_accounts(ldb,
 309                                                 opts.instance_id,
 310                                                 len(conversations),
 311                                                 opts.fixed_password)
 312
 313     statsdir = traffic.mk_masked_dir(tempdir, 'stats')
 314
 315     if opts.traffic_summary:
 316         if opts.traffic_summary == '-':
 317             summary_dest = sys.stdout
 318         else:
 319             summary_dest = open(opts.traffic_summary, 'w')
 320
 321         print >>sys.stderr, "Writing traffic summary"
 322         summaries = []
 323         for c in conversations:
 324             summaries += c.replay_as_summary_lines()
 325
 326         summaries.sort()
 327         for (time, line) in summaries:
 328             print >>summary_dest, line
 329
 330         exit(0)
 331
 332     traffic.replay(conversations, host,
 333                    lp=lp,
 334                    creds=creds,
 335                    accounts=accounts,
 336                    dns_rate=opts.dns_rate,
 337                    duration=duration,
 338                    badpassword_frequency=opts.badpassword_frequency,
 339                    prefer_kerberos=opts.prefer_kerberos,
 340                    statsdir=statsdir,
 341                    domain=domain,
 342                    base_dn=ldb.domain_dn(),
 343                    ou=traffic.ou_name(ldb, opts.instance_id),
 344                    tempdir=tempdir,
 345                    domain_sid=ldb.get_domain_sid())
 346
 347     if opts.timing_data == '-':
 348         timing_dest = sys.stdout
 349     elif opts.timing_data is None:
 350         timing_dest = None
 351     else:
 352         timing_dest = open(opts.timing_data, 'w')
 353
 354     print >>sys.stderr, "Generating statistics"
 355     traffic.generate_stats(statsdir, timing_dest)
 356
 357     if not opts.preserve_tempdir:
 358         print >>sys.stderr, "Removing temporary directory"
 359         shutil.rmtree(tempdir)
 360
 361
 362 main()