totp-generator/totp.py

#! /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()