diff options
-rw-r--r-- | dnsviz/commands/graph.py | 6 | ||||
-rw-r--r-- | dnsviz/commands/grok.py | 6 | ||||
-rw-r--r-- | dnsviz/commands/print.py | 6 | ||||
-rw-r--r-- | dnsviz/util.py | 10 |
4 files changed, 19 insertions, 9 deletions
diff --git a/dnsviz/commands/graph.py b/dnsviz/commands/graph.py index 2719b1d..7d3fe06 100644 --- a/dnsviz/commands/graph.py +++ b/dnsviz/commands/graph.py @@ -44,7 +44,7 @@ import dns.exception, dns.name from dnsviz.analysis import OfflineDomainNameAnalysis, DNS_RAW_VERSION from dnsviz.config import DNSVIZ_SHARE_PATH, JQUERY_PATH, JQUERY_UI_PATH, JQUERY_UI_CSS_PATH, RAPHAEL_PATH from dnsviz.format import latin1_binary_to_string as lb2s -from dnsviz.util import get_trusted_keys, get_default_trusted_keys +from dnsviz.util import get_trusted_keys, get_default_trusted_keys, io_try_buffered # If the import of DNSAuthGraph fails because of the lack of pygraphviz, it # will be reported later @@ -154,8 +154,8 @@ class GraphArgHelper: self.parser = argparse.ArgumentParser(description='Graph the assessment of diagnostic DNS queries', prog=prog) # python3/python2 dual compatibility - stdin_buffer = io.open(sys.stdin.fileno(), 'rb', closefd=False) - stdout_buffer = io.open(sys.stdout.fileno(), 'wb', closefd=False) + stdin_buffer = io_try_buffered(sys.stdin, 'rb', closefd=False) + stdout_buffer = io_try_buffered(sys.stdout, 'wb', closefd=False) try: self.parser.add_argument('-f', '--names-file', diff --git a/dnsviz/commands/grok.py b/dnsviz/commands/grok.py index 6f2d8cd..8fc77c3 100644 --- a/dnsviz/commands/grok.py +++ b/dnsviz/commands/grok.py @@ -43,7 +43,7 @@ import dns.exception, dns.name from dnsviz.analysis import OfflineDomainNameAnalysis, DNS_RAW_VERSION from dnsviz.format import latin1_binary_to_string as lb2s -from dnsviz.util import get_trusted_keys +from dnsviz.util import get_trusted_keys, io_try_buffered # If the import of DNSAuthGraph fails because of the lack of pygraphviz, it # will be reported later @@ -170,8 +170,8 @@ class GrokArgHelper: self.parser = argparse.ArgumentParser(description='Assess diagnostic DNS queries', prog=prog) # python3/python2 dual compatibility - stdin_buffer = io.open(sys.stdin.fileno(), 'rb', closefd=False) - stdout_buffer = io.open(sys.stdout.fileno(), 'wb', closefd=False) + stdin_buffer = io_try_buffered(sys.stdin, 'rb', closefd=False) + stdout_buffer = io_try_buffered(sys.stdout, 'wb', closefd=False) try: self.parser.add_argument('-f', '--names-file', diff --git a/dnsviz/commands/print.py b/dnsviz/commands/print.py index 3dcbb8d..80aa4a5 100644 --- a/dnsviz/commands/print.py +++ b/dnsviz/commands/print.py @@ -43,7 +43,7 @@ import dns.exception, dns.name from dnsviz.analysis import TTLAgnosticOfflineDomainNameAnalysis, DNS_RAW_VERSION from dnsviz.format import latin1_binary_to_string as lb2s -from dnsviz.util import get_trusted_keys, get_default_trusted_keys +from dnsviz.util import get_trusted_keys, get_default_trusted_keys, io_try_buffered # If the import of DNSAuthGraph fails because of the lack of pygraphviz, it # will be reported later @@ -312,8 +312,8 @@ class PrintArgHelper: self.parser = argparse.ArgumentParser(description='Print the assessment of diagnostic DNS queries', prog=prog) # python3/python2 dual compatibility - stdin_buffer = io.open(sys.stdin.fileno(), 'rb', closefd=False) - stdout_buffer = io.open(sys.stdout.fileno(), 'wb', closefd=False) + stdin_buffer = io_try_buffered(sys.stdin, 'rb', closefd=False) + stdout_buffer = io_try_buffered(sys.stdout, 'wb', closefd=False) try: self.parser.add_argument('-f', '--names-file', diff --git a/dnsviz/util.py b/dnsviz/util.py index 8efed82..f4c6c71 100644 --- a/dnsviz/util.py +++ b/dnsviz/util.py @@ -184,3 +184,13 @@ def get_client_address(server): ip = IPAddr(s.getsockname()[0]) s.close() return ip + +def io_try_buffered(f, mode, closefd=True): + """Try opening buffered reader, but allow unbuffered on failure. + + Required to pass tests under pytest.""" + try: + return io.open(f.fileno(), mode, closefd=closefd) + except io.UnsupportedOperation: + # raised by f.fileno() + return f |