diff --git a/extract_otp_secret_keys.py b/extract_otp_secret_keys.py index 51d3e6d..5dfa313 100644 --- a/extract_otp_secret_keys.py +++ b/extract_otp_secret_keys.py @@ -140,12 +140,15 @@ def extract_otps(args): def get_lines_from_file(filename): lines = read_lines_from_text_file(filename) - if lines != None: + if are_bytes(lines): + abort('\nBinary input was given in stdin, please use = instead of -.') + elif lines: return lines # could not process text file, try reading as image return convert_img_to_line(filename) + def read_lines_from_text_file(filename): if filename != '=': check_file_exists(filename) @@ -169,6 +172,7 @@ def read_lines_from_text_file(filename): finally: finput.close() + def convert_img_to_line(filename): if filename != '-': try: @@ -180,7 +184,10 @@ def convert_img_to_line(filename): except AttributeError: # Workaround for pytest, since pytest cannot monkeypatch sys.stdin.buffer stdin = sys.stdin.read() - array = frombuffer(stdin, dtype='uint8') + try: + array = frombuffer(stdin, dtype='uint8') + except TypeError as e: + abort('\nERROR: Cannot read binary stdin buffer. Exception: {}'.format(str(e))) image = imdecode(array, IMREAD_UNCHANGED) if image is None: @@ -384,6 +391,16 @@ def check_file_exists(filename): abort('\nERROR: Input file provided is non-existent or not a file.' '\ninput file: {}'.format(filename)) + +def are_bytes(lines): + if lines and len(lines) > 0: + try: + lines[0].startswith('#') + return False + except (UnicodeDecodeError, AttributeError, TypeError): + return True + + def eprint(*args, **kwargs): '''Print to stderr.''' print(*args, file=sys.stderr, **kwargs) diff --git a/test_extract_otp_secret_keys_pytest.py b/test_extract_otp_secret_keys_pytest.py index 99d8bf6..8d0e561 100644 --- a/test_extract_otp_secret_keys_pytest.py +++ b/test_extract_otp_secret_keys_pytest.py @@ -68,6 +68,25 @@ def test_extract_stdin_stdout(capsys, monkeypatch): assert captured.err == '' +def test_extract_stdin_stdout_wrong_symbol(capsys, monkeypatch): + # Arrange + monkeypatch.setattr('sys.stdin', StringIO(read_file_to_str('example_export.txt'))) + + # Act + with raises(SystemExit) as e: + extract_otp_secret_keys.main(['=']) + + # Assert + captured = capsys.readouterr() + + expected_stderr = "\nERROR: Cannot read binary stdin buffer. Exception: a bytes-like object is required, not 'str'\n" + + assert captured.err == expected_stderr + assert captured.out == '' + assert e.value.code == 1 + assert e.type == SystemExit + + def test_extract_csv(capsys): # Arrange cleanup() @@ -534,6 +553,26 @@ Type: totp assert captured.err == '' +def test_img_qr_reader_from_stdin_wrong_symbol(capsys, monkeypatch): + # Arrange + # sys.stdin.buffer should be monkey patched, but it does not work + monkeypatch.setattr('sys.stdin', read_binary_file_as_stream('test/test_googleauth_export.png')) + + # Act + with raises(SystemExit) as e: + extract_otp_secret_keys.main(['-']) + + # Assert + captured = capsys.readouterr() + + expected_stderr = '\nBinary input was given in stdin, please use = instead of -.\n' + + assert captured.err == expected_stderr + assert captured.out == '' + assert e.value.code == 1 + assert e.type == SystemExit + + def test_img_qr_reader_no_qr_code_in_image(capsys): # Act with raises(SystemExit) as e: