scripts: Scripts to replay and generate samba traffic
[garming/samba-autobuild/.git] / 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()