mirror of
https://github.com/scito/extract_otp_secret_keys
synced 2024-11-18 09:28:19 +00:00
organize imports
This commit is contained in:
scito
parent
b0b4c29e7b
commit
2bcaa35251
@ -41,18 +41,19 @@
|
||||
# You should have received a copy of the GNU General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
# TODO optimze imports
|
||||
import argparse
|
||||
import base64
|
||||
import fileinput
|
||||
import sys
|
||||
import csv
|
||||
import fileinput
|
||||
import json
|
||||
from cv2 import imread, imdecode, IMREAD_UNCHANGED
|
||||
from urllib.parse import parse_qs, urlencode, urlparse, quote
|
||||
from os import path, makedirs
|
||||
from re import compile as rcompile
|
||||
from numpy import frombuffer
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import urllib.parse as urlparse
|
||||
|
||||
import cv2
|
||||
import numpy
|
||||
|
||||
import protobuf_generated_python.google_auth_pb2
|
||||
|
||||
|
||||
@ -79,7 +80,7 @@ def main(sys_args):
|
||||
def parse_args(sys_args):
|
||||
formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=52)
|
||||
arg_parser = argparse.ArgumentParser(formatter_class=formatter)
|
||||
arg_parser.add_argument('infile', help='1) file or - for stdin with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored; 2) image file containing a QR code or = for stdin for an image containing a QR code', nargs='+')
|
||||
arg_parser.add_argument('infile', help='1) file or - for stdin with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored; or 2) image file containing a QR code or = for stdin for an image containing a QR code', nargs='+')
|
||||
arg_parser.add_argument('--json', '-j', help='export json file or - for stdout', metavar=('FILE'))
|
||||
arg_parser.add_argument('--csv', '-c', help='export csv file or - for stdout', metavar=('FILE'))
|
||||
arg_parser.add_argument('--keepass', '-k', help='export totp/hotp csv file(s) for KeePass, - for stdout', metavar=('FILE'))
|
||||
@ -180,7 +181,7 @@ def convert_img_to_line(filename):
|
||||
if verbose: print('Reading image {}'.format(filename))
|
||||
try:
|
||||
if filename != '=':
|
||||
image = imread(filename)
|
||||
image = cv2.imread(filename)
|
||||
else:
|
||||
try:
|
||||
stdin = sys.stdin.buffer.read()
|
||||
@ -188,10 +189,10 @@ def convert_img_to_line(filename):
|
||||
# Workaround for pytest, since pytest cannot monkeypatch sys.stdin.buffer
|
||||
stdin = sys.stdin.read()
|
||||
try:
|
||||
array = frombuffer(stdin, dtype='uint8')
|
||||
array = numpy.frombuffer(stdin, dtype='uint8')
|
||||
except TypeError as e:
|
||||
abort('\nERROR: Cannot read binary stdin buffer. Exception: {}'.format(str(e)))
|
||||
image = imdecode(array, IMREAD_UNCHANGED)
|
||||
image = cv2.imdecode(array, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
if image is None:
|
||||
abort('\nERROR: Unable to open file for reading.\ninput file: {}'.format(filename))
|
||||
@ -220,10 +221,10 @@ def get_payload_from_line(line, i, infile):
|
||||
global verbose
|
||||
if not line.startswith('otpauth-migration://'):
|
||||
eprint( '\nWARN: line is not a otpauth-migration:// URL\ninput file: {}\nline "{}"\nProbably a wrong file was given'.format(infile, line))
|
||||
parsed_url = urlparse(line)
|
||||
parsed_url = urlparse.urlparse(line)
|
||||
if verbose > 1: print('\nDEBUG: parsed_url={}'.format(parsed_url))
|
||||
try:
|
||||
params = parse_qs(parsed_url.query, strict_parsing=True)
|
||||
params = urlparse.parse_qs(parsed_url.query, strict_parsing=True)
|
||||
except: # Not necessary for Python >= 3.11
|
||||
params = []
|
||||
if verbose > 1: print('\nDEBUG: querystring params={}'.format(params))
|
||||
@ -264,7 +265,7 @@ def build_otp_url(secret, raw_otp):
|
||||
url_params = {'secret': secret}
|
||||
if raw_otp.type == 1: url_params['counter'] = raw_otp.counter
|
||||
if raw_otp.issuer: url_params['issuer'] = raw_otp.issuer
|
||||
otp_url = 'otpauth://{}/{}?'.format(get_otp_type_str_from_code(raw_otp.type), quote(raw_otp.name)) + urlencode( url_params)
|
||||
otp_url = 'otpauth://{}/{}?'.format(get_otp_type_str_from_code(raw_otp.type), urlparse.quote(raw_otp.name)) + urlparse.urlencode( url_params)
|
||||
return otp_url
|
||||
|
||||
|
||||
@ -281,8 +282,8 @@ def print_otp(otp):
|
||||
|
||||
def save_qr(otp, args, j):
|
||||
dir = args.saveqr
|
||||
if not (path.exists(dir)): makedirs(dir, exist_ok=True)
|
||||
pattern = rcompile(r'[\W_]+')
|
||||
if not (os.path.exists(dir)): os.makedirs(dir, exist_ok=True)
|
||||
pattern = re.compile(r'[\W_]+')
|
||||
file_otp_name = pattern.sub('', otp['name'])
|
||||
file_otp_issuer = pattern.sub('', otp['issuer'])
|
||||
save_qr_file(args, otp['url'], '{}/{}-{}{}.png'.format(dir, j, file_otp_name, '-' + file_otp_issuer if file_otp_issuer else ''))
|
||||
@ -374,7 +375,7 @@ def has_otp_type(otps, otp_type):
|
||||
|
||||
def add_pre_suffix(file, pre_suffix):
|
||||
'''filename.ext, pre -> filename.pre.ext'''
|
||||
name, ext = path.splitext(file)
|
||||
name, ext = os.path.splitext(file)
|
||||
return name + "." + pre_suffix + (ext if ext else "")
|
||||
|
||||
|
||||
@ -394,7 +395,7 @@ def open_file_or_stdout_for_csv(filename):
|
||||
|
||||
|
||||
def check_file_exists(filename):
|
||||
if filename != '-' and not path.isfile(filename):
|
||||
if filename != '-' and not os.path.isfile(filename):
|
||||
abort('\nERROR: Input file provided is non-existent or not a file.'
|
||||
'\ninput file: {}'.format(filename))
|
||||
|
||||
|
||||
@ -18,13 +18,16 @@
|
||||
# You should have received a copy of the GNU General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
from utils import read_csv, read_csv_str, read_json, read_json_str, remove_files, remove_dir_with_files, read_file_to_str, read_binary_file_as_stream, file_exits
|
||||
from os import path
|
||||
from pytest import raises, mark
|
||||
from io import StringIO, BytesIO
|
||||
from sys import implementation, stdin
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
|
||||
from pytest import mark, raises
|
||||
|
||||
import extract_otp_secret_keys
|
||||
from utils import (file_exits, read_binary_file_as_stream, read_csv,
|
||||
read_csv_str, read_file_to_str, read_json, read_json_str,
|
||||
remove_dir_with_files, remove_files)
|
||||
|
||||
|
||||
def test_extract_stdout(capsys):
|
||||
@ -71,7 +74,7 @@ def test_extract_non_existent_file(capsys):
|
||||
|
||||
def test_extract_stdin_stdout(capsys, monkeypatch):
|
||||
# Arrange
|
||||
monkeypatch.setattr('sys.stdin', StringIO(read_file_to_str('example_export.txt')))
|
||||
monkeypatch.setattr('sys.stdin', io.StringIO(read_file_to_str('example_export.txt')))
|
||||
|
||||
# Act
|
||||
extract_otp_secret_keys.main(['-'])
|
||||
@ -85,7 +88,7 @@ def test_extract_stdin_stdout(capsys, monkeypatch):
|
||||
|
||||
def test_extract_stdin_stdout_wrong_symbol(capsys, monkeypatch):
|
||||
# Arrange
|
||||
monkeypatch.setattr('sys.stdin', StringIO(read_file_to_str('example_export.txt')))
|
||||
monkeypatch.setattr('sys.stdin', io.StringIO(read_file_to_str('example_export.txt')))
|
||||
|
||||
# Act
|
||||
with raises(SystemExit) as e:
|
||||
@ -149,7 +152,7 @@ def test_extract_csv_stdout(capsys):
|
||||
def test_extract_stdin_and_csv_stdout(capsys, monkeypatch):
|
||||
# Arrange
|
||||
cleanup()
|
||||
monkeypatch.setattr('sys.stdin', StringIO(read_file_to_str('example_export.txt')))
|
||||
monkeypatch.setattr('sys.stdin', io.StringIO(read_file_to_str('example_export.txt')))
|
||||
|
||||
# Act
|
||||
extract_otp_secret_keys.main(['-c', '-', '-'])
|
||||
@ -347,16 +350,16 @@ def test_extract_saveqr(capsys):
|
||||
assert captured.out == ''
|
||||
assert captured.err == ''
|
||||
|
||||
assert path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png')
|
||||
assert path.isfile('testout/qr/2-piraspberrypi.png')
|
||||
assert path.isfile('testout/qr/3-piraspberrypi.png')
|
||||
assert path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png')
|
||||
assert os.path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png')
|
||||
assert os.path.isfile('testout/qr/2-piraspberrypi.png')
|
||||
assert os.path.isfile('testout/qr/3-piraspberrypi.png')
|
||||
assert os.path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png')
|
||||
|
||||
# Clean up
|
||||
cleanup()
|
||||
|
||||
|
||||
@mark.skipif(implementation.name == 'pypy', reason="Encoding problems in verbose mode in pypy.")
|
||||
@mark.skipif(sys.implementation.name == 'pypy', reason="Encoding problems in verbose mode in pypy.")
|
||||
def test_extract_verbose(capsys):
|
||||
# Act
|
||||
extract_otp_secret_keys.main(['-v', 'example_export.txt'])
|
||||
|
||||
@ -22,8 +22,8 @@ import unittest
|
||||
import io
|
||||
from contextlib import redirect_stdout
|
||||
from utils import read_csv, read_json, remove_file, remove_dir_with_files, Capturing, read_file_to_str
|
||||
from os import path
|
||||
from sys import implementation
|
||||
import os
|
||||
import sys
|
||||
|
||||
import extract_otp_secret_keys
|
||||
|
||||
@ -160,13 +160,13 @@ Type: totp
|
||||
def test_extract_saveqr(self):
|
||||
extract_otp_secret_keys.main(['-q', '-s', 'testout/qr/', 'example_export.txt'])
|
||||
|
||||
self.assertTrue(path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png'))
|
||||
self.assertTrue(path.isfile('testout/qr/2-piraspberrypi.png'))
|
||||
self.assertTrue(path.isfile('testout/qr/3-piraspberrypi.png'))
|
||||
self.assertTrue(path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png'))
|
||||
self.assertTrue(os.path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png'))
|
||||
self.assertTrue(os.path.isfile('testout/qr/2-piraspberrypi.png'))
|
||||
self.assertTrue(os.path.isfile('testout/qr/3-piraspberrypi.png'))
|
||||
self.assertTrue(os.path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png'))
|
||||
|
||||
def test_extract_verbose(self):
|
||||
if implementation.name == 'pypy': self.skipTest("Encoding problems in verbose mode in pypy.")
|
||||
if sys.implementation.name == 'pypy': self.skipTest("Encoding problems in verbose mode in pypy.")
|
||||
out = io.StringIO()
|
||||
with redirect_stdout(out):
|
||||
extract_otp_secret_keys.main(['-v', 'example_export.txt'])
|
||||
|
||||
8
utils.py
8
utils.py
@ -17,7 +17,7 @@ import csv
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from io import StringIO, BytesIO
|
||||
import io
|
||||
import sys
|
||||
import glob
|
||||
|
||||
@ -31,9 +31,9 @@ with Capturing() as output:
|
||||
'''
|
||||
def __enter__(self):
|
||||
self._stdout = sys.stdout
|
||||
sys.stdout = self._stringio_std = StringIO()
|
||||
sys.stdout = self._stringio_std = io.StringIO()
|
||||
self._stderr = sys.stderr
|
||||
sys.stderr = self._stringio_err = StringIO()
|
||||
sys.stderr = self._stringio_err = io.StringIO()
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
@ -106,4 +106,4 @@ def read_file_to_str(filename):
|
||||
def read_binary_file_as_stream(filename):
|
||||
"""Returns binary file content."""
|
||||
with open(filename, "rb",) as infile:
|
||||
return BytesIO(infile.read())
|
||||
return io.BytesIO(infile.read())
|
||||
|
||||
Loading…
Reference in New Issue
Block a user