totp-generator/totp.py

141 lines
3.6 KiB
Python
Raw Permalink Normal View History

2022-06-25 13:38:39 +02:00
#! /usr/bin/python3
import cv2
import pyzbar.pyzbar as pyzbar
import re
import qrcode
from pathlib import Path
import json
from typing import Union, Dict
import urllib.parse
ROOT_DIR = Path(__file__).resolve().parent
QRCODE_INPUT_DIR = ROOT_DIR / 'in'
JSON_DIR = ROOT_DIR / 'json'
QRCODE_OUTPUT_DIR = ROOT_DIR / 'out'
for d in [QRCODE_INPUT_DIR, JSON_DIR, QRCODE_OUTPUT_DIR]:
d.mkdir(exist_ok=True)
def normalize(string: str):
return string.strip().lower().replace(' ', '_')
def make_qr_code(data):
qr = qrcode.QRCode(
error_correction=qrcode.constants.ERROR_CORRECT_H,
version=None,
box_size=12,
border=5,
)
qr.add_data(str(data))
qr.make(fit=True)
img = qr.make_image()
return img
def make_totp_url(secret, issuer='XXX', username='xxx', **kwargs):
return "otpauth://totp/{issuer}:{username}?secret={secret}&issuer={issuer}".format(
issuer=issuer, secret=secret, username=username
)
def read_qr_code(path):
img = cv2.imread(str(path))
dec = pyzbar.decode(img, symbols=[pyzbar.ZBarSymbol.QRCODE])
for obj in dec:
return obj.data.decode('ascii')
def parse_totp_link(string) -> Union[Dict, None]:
match = re.search(
r'otpauth://totp/'
r'(?:(?P<issuer>.*):)?(?P<username>.*?)?'
r'\?secret=(?P<secret>.*?)'
r'(?:'
r'(&digits=(?P<digits>\d+?))|'
r'(&algorithm=(?P<algorithm>.+?))|'
r'(&issuer=(?P<issuer2>.+?))|'
r'(&period=(?P<period>\d+?))|'
r')*$',
string
)
if not match:
return None
else:
d = match.groupdict()
try:
if d['issuer'] and d['issuer2'] and d['issuer'] != d['issuer2']:
print('different issuers: confused')
exit(1)
except KeyError:
pass
try:
if d['issuer2']:
d['issuer'] = d['issuer2']
except KeyError:
pass
try:
d.pop('issuer2')
except KeyError:
pass
d['url'] = string
return d
def parse_qr_codes():
all = []
for qr in QRCODE_INPUT_DIR.rglob('*'):
print(f'reading in: {qr.name}' )
link = urllib.parse.unquote(read_qr_code(qr))
if not link:
continue
totp = parse_totp_link(link)
try:
if not totp['issuer'] == totp['issuer2']:
print('different issuers found. confused')
else:
totp.pop('issuer2')
except KeyError:
pass
if totp is not None:
all.append(totp)
with open(JSON_DIR / qr.with_suffix('.json').name, 'w') as f:
json.dump(totp, f, indent=4, sort_keys=True)
with open(JSON_DIR / 'all.json', 'w') as f:
json.dump(all, f, indent=4, sort_keys=True)
def generate_qr_code(totp):
img = make_qr_code(make_totp_url(**totp))
filename = '{issuer}:{username}.png'.format(
issuer=totp['issuer'] if 'issuer' in totp.keys() else 'XXX',
username=totp['username'] if 'username' in totp.keys() else 'xxx'
)
img.save(QRCODE_OUTPUT_DIR / normalize(filename))
def generate_qr_codes():
for js in JSON_DIR.rglob('*.json'):
if js.name == 'all.json':
continue
with open(js, 'r') as f:
totp = json.load(f)
generate_qr_code(totp)
def produce_from_all():
with open(JSON_DIR / 'all.json') as f:
l = json.load(f)
l = list(l)
for d in l:
generate_qr_code(d)
if __name__ == "__main__":
generate_qr_codes()
parse_qr_codes()