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