Take advantage of the automatic file handling thinfy in argparser
This commit is contained in:
parent
e6b3345421
commit
f390da248c
|
@ -34,10 +34,13 @@ def main():
|
||||||
|
|
||||||
# -o can be used to write the .sshwot thing to a file instead of stdout
|
# -o can be used to write the .sshwot thing to a file instead of stdout
|
||||||
parser.add_argument('-o',
|
parser.add_argument('-o',
|
||||||
# We store one argument given after this one to the property
|
# Given one argument, we open a file of that name and store it
|
||||||
# outfile, which will be None otherwise
|
# to outfile, which will be sys.stdout.buffer otherwise.
|
||||||
|
# We use .buffer since we're going to write binary data
|
||||||
action = 'store',
|
action = 'store',
|
||||||
dest = 'outfile',
|
dest = 'outfile',
|
||||||
|
type = argparse.FileType('wb'),
|
||||||
|
default = sys.stdout.buffer,
|
||||||
# This is what will be displayed in the help after -o
|
# This is what will be displayed in the help after -o
|
||||||
metavar = 'outfile',
|
metavar = 'outfile',
|
||||||
help = 'write the sshwot file to a given file instead of the stdout'
|
help = 'write the sshwot file to a given file instead of the stdout'
|
||||||
|
@ -68,15 +71,7 @@ def main():
|
||||||
for known_hosts_entry in known_host_entries:
|
for known_hosts_entry in known_host_entries:
|
||||||
entries.append(process_known_hosts.known_hosts_to_entry(known_hosts_entry))
|
entries.append(process_known_hosts.known_hosts_to_entry(known_hosts_entry))
|
||||||
|
|
||||||
if args.outfile is None:
|
write_file.write(args.outfile, entries)
|
||||||
# Write to stdout by default
|
|
||||||
# We use sys.stdout.buffer instead of just sys.stdout because we
|
|
||||||
# are writing binary data
|
|
||||||
write_file.write(sys.stdout.buffer, entries)
|
|
||||||
|
|
||||||
else:
|
|
||||||
with open(args.outfile, 'wb') as outf:
|
|
||||||
write_file.write(outf, entries)
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
if __name__ == '__main__':
|
||||||
try:
|
try:
|
||||||
|
|
Loading…
Reference in New Issue