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