smokephil/extract_otp_secret_keys
1
0
Fork 0
mirror of https://github.com/scito/extract_otp_secret_keys.git synced 2025-12-08 15:54:58 +01: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
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 # You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>. # along with this program. If not, see <https://www.gnu.org/licenses/>.
# TODO optimze imports
import argparse import argparse
import base64 import base64
import fileinput
import sys
import csv import csv
import fileinput
import json import json
from cv2 import imread, imdecode, IMREAD_UNCHANGED import os
from urllib.parse import parse_qs, urlencode, urlparse, quote import re
from os import path, makedirs import sys
from re import compile as rcompile import urllib.parse as urlparse
from numpy import frombuffer
import cv2
import numpy
import protobuf_generated_python.google_auth_pb2 import protobuf_generated_python.google_auth_pb2
@ -79,7 +80,7 @@ def main(sys_args):
def parse_args(sys_args): def parse_args(sys_args):
formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=52) formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=52)
arg_parser = argparse.ArgumentParser(formatter_class=formatter) 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('--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('--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')) 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)) if verbose: print('Reading image {}'.format(filename))
try: try:
if filename != '=': if filename != '=':
image = imread(filename) image = cv2.imread(filename)
else: else:
try: try:
stdin = sys.stdin.buffer.read() 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 # Workaround for pytest, since pytest cannot monkeypatch sys.stdin.buffer
stdin = sys.stdin.read() stdin = sys.stdin.read()
try: try:
array = frombuffer(stdin, dtype='uint8') array = numpy.frombuffer(stdin, dtype='uint8')
except TypeError as e: except TypeError as e:
abort('\nERROR: Cannot read binary stdin buffer. Exception: {}'.format(str(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: if image is None:
abort('\nERROR: Unable to open file for reading.\ninput file: {}'.format(filename)) 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 global verbose
if not line.startswith('otpauth-migration://'): 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)) 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)) if verbose > 1: print('\nDEBUG: parsed_url={}'.format(parsed_url))
try: 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 except: # Not necessary for Python >= 3.11
params = [] params = []
if verbose > 1: print('\nDEBUG: querystring params={}'.format(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} url_params = {'secret': secret}
if raw_otp.type == 1: url_params['counter'] = raw_otp.counter if raw_otp.type == 1: url_params['counter'] = raw_otp.counter
if raw_otp.issuer: url_params['issuer'] = raw_otp.issuer 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 return otp_url
@ -281,8 +282,8 @@ def print_otp(otp):
def save_qr(otp, args, j): def save_qr(otp, args, j):
dir = args.saveqr dir = args.saveqr
if not (path.exists(dir)): makedirs(dir, exist_ok=True) if not (os.path.exists(dir)): os.makedirs(dir, exist_ok=True)
pattern = rcompile(r'[\W_]+') pattern = re.compile(r'[\W_]+')
file_otp_name = pattern.sub('', otp['name']) file_otp_name = pattern.sub('', otp['name'])
file_otp_issuer = pattern.sub('', otp['issuer']) 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 '')) 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): def add_pre_suffix(file, pre_suffix):
'''filename.ext, pre -> filename.pre.ext''' '''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 "") 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): 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.' abort('\nERROR: Input file provided is non-existent or not a file.'
'\ninput file: {}'.format(filename)) '\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 # You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>. # 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 import io
from os import path import os
from pytest import raises, mark import sys
from io import StringIO, BytesIO
from sys import implementation, stdin from pytest import mark, raises
import extract_otp_secret_keys 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): def test_extract_stdout(capsys):
@ -71,7 +74,7 @@ def test_extract_non_existent_file(capsys):
def test_extract_stdin_stdout(capsys, monkeypatch): def test_extract_stdin_stdout(capsys, monkeypatch):
# Arrange # 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 # Act
extract_otp_secret_keys.main(['-']) 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): def test_extract_stdin_stdout_wrong_symbol(capsys, monkeypatch):
# Arrange # 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 # Act
with raises(SystemExit) as e: with raises(SystemExit) as e:
@ -149,7 +152,7 @@ def test_extract_csv_stdout(capsys):
def test_extract_stdin_and_csv_stdout(capsys, monkeypatch): def test_extract_stdin_and_csv_stdout(capsys, monkeypatch):
# Arrange # Arrange
cleanup() 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 # Act
extract_otp_secret_keys.main(['-c', '-', '-']) extract_otp_secret_keys.main(['-c', '-', '-'])
@ -347,16 +350,16 @@ def test_extract_saveqr(capsys):
assert captured.out == '' assert captured.out == ''
assert captured.err == '' assert captured.err == ''
assert path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png') assert os.path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png')
assert path.isfile('testout/qr/2-piraspberrypi.png') assert os.path.isfile('testout/qr/2-piraspberrypi.png')
assert path.isfile('testout/qr/3-piraspberrypi.png') assert os.path.isfile('testout/qr/3-piraspberrypi.png')
assert path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png') assert os.path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png')
# Clean up # Clean up
cleanup() 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): def test_extract_verbose(capsys):
# Act # Act
extract_otp_secret_keys.main(['-v', 'example_export.txt']) 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 import io
from contextlib import redirect_stdout from contextlib import redirect_stdout
from utils import read_csv, read_json, remove_file, remove_dir_with_files, Capturing, read_file_to_str from utils import read_csv, read_json, remove_file, remove_dir_with_files, Capturing, read_file_to_str
from os import path import os
from sys import implementation import sys
import extract_otp_secret_keys import extract_otp_secret_keys
@ -160,13 +160,13 @@ Type: totp
def test_extract_saveqr(self): def test_extract_saveqr(self):
extract_otp_secret_keys.main(['-q', '-s', 'testout/qr/', 'example_export.txt']) extract_otp_secret_keys.main(['-q', '-s', 'testout/qr/', 'example_export.txt'])
self.assertTrue(path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png')) self.assertTrue(os.path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png'))
self.assertTrue(path.isfile('testout/qr/2-piraspberrypi.png')) self.assertTrue(os.path.isfile('testout/qr/2-piraspberrypi.png'))
self.assertTrue(path.isfile('testout/qr/3-piraspberrypi.png')) self.assertTrue(os.path.isfile('testout/qr/3-piraspberrypi.png'))
self.assertTrue(path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png')) self.assertTrue(os.path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png'))
def test_extract_verbose(self): 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() out = io.StringIO()
with redirect_stdout(out): with redirect_stdout(out):
extract_otp_secret_keys.main(['-v', 'example_export.txt']) extract_otp_secret_keys.main(['-v', 'example_export.txt'])

8
utils.py
View file

@ -17,7 +17,7 @@ import csv
import json import json
import os import os
import shutil import shutil
from io import StringIO, BytesIO import io
import sys import sys
import glob import glob
@ -31,9 +31,9 @@ with Capturing() as output:
''' '''
def __enter__(self): def __enter__(self):
self._stdout = sys.stdout self._stdout = sys.stdout
sys.stdout = self._stringio_std = StringIO() sys.stdout = self._stringio_std = io.StringIO()
self._stderr = sys.stderr self._stderr = sys.stderr
sys.stderr = self._stringio_err = StringIO() sys.stderr = self._stringio_err = io.StringIO()
return self return self
def __exit__(self, *args): def __exit__(self, *args):
@ -106,4 +106,4 @@ def read_file_to_str(filename):
def read_binary_file_as_stream(filename): def read_binary_file_as_stream(filename):
"""Returns binary file content.""" """Returns binary file content."""
with open(filename, "rb",) as infile: with open(filename, "rb",) as infile:
return BytesIO(infile.read()) return io.BytesIO(infile.read())