diff --git a/src/main-export-known-hosts.py b/src/main-export-known-hosts.py index 86d3185..63288b9 100644 --- a/src/main-export-known-hosts.py +++ b/src/main-export-known-hosts.py @@ -34,10 +34,13 @@ def main(): # -o can be used to write the .sshwot thing to a file instead of stdout parser.add_argument('-o', - # We store one argument given after this one to the property - # outfile, which will be None otherwise + # Given one argument, we open a file of that name and store it + # to outfile, which will be sys.stdout.buffer otherwise. + # We use .buffer since we're going to write binary data action = 'store', dest = 'outfile', + type = argparse.FileType('wb'), + default = sys.stdout.buffer, # This is what will be displayed in the help after -o metavar = 'outfile', 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: entries.append(process_known_hosts.known_hosts_to_entry(known_hosts_entry)) - if args.outfile is None: - # 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) + write_file.write(args.outfile, entries) if __name__ == '__main__': try: