Archives/extract_otp_secret_keys
2
0
Fork 0
mirror of https://github.com/scito/extract_otp_secret_keys synced 2024-11-04 12:00:14 +00:00
Code Issues Projects Releases Wiki Activity

organize imports

Browse Source
This commit is contained in:
scito 2022-12-25 11:00:15 +01:00
parent b0b4c29e7b
commit 2bcaa35251
4 changed files with 47 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
extract_otp_secret_keys.py
View File

@ -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))

29
test_extract_otp_secret_keys_pytest.py
View File

@ -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'])

14
test_extract_otp_secret_keys_unittest.py
View File

@ -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
View File

@ -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())