Archives/extract_otp_secret_keys
2
0
Fork 0
mirror of https://github.com/scito/extract_otp_secret_keys synced 2024-11-16 06:15:17 +00:00
Code Issues Projects Releases Wiki Activity

extract draw_box and print_text functions

Browse Source
This commit is contained in:
scito 2023-01-01 19:22:46 +01:00
parent 16047a5b15
commit 8545dab7a5
1 changed files with 44 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

75
src/extract_otp_secrets.py
View File

@ -53,7 +53,6 @@ import re
import sys import sys
import urllib.parse as urlparse import urllib.parse as urlparse
from enum import Enum from enum import Enum
from operator import add
from typing import Any, List, Optional, TextIO, Tuple, Union from typing import Any, List, Optional, TextIO, Tuple, Union
# workaround for PYTHON <= 3.7: compatibility # workaround for PYTHON <= 3.7: compatibility
@ -68,8 +67,9 @@ import protobuf_generated_python.google_auth_pb2 as pb
import colorama import colorama
try: try:
import cv2 # type: ignore import cv2 # type: ignore # TODO use cv2 types if available
import numpy
import numpy as np # TODO use numpy types if available
try: try:
import pyzbar.pyzbar as zbar # type: ignore import pyzbar.pyzbar as zbar # type: ignore
@ -91,13 +91,16 @@ Exception: {e}""")
FONT_SCALE: Final[int] = 1 FONT_SCALE: Final[int] = 1
FONT_THICKNESS: Final[int] = 1 FONT_THICKNESS: Final[int] = 1
FONT_LINE_STYLE: Final[int] = cv2.LINE_AA FONT_LINE_STYLE: Final[int] = cv2.LINE_AA
RECT_THICKNESS: Final[int] = 5 BOX_THICKNESS: Final[int] = 5
# workaround for PYTHON <= 3.7: must use () for assignments # workaround for PYTHON <= 3.7: must use () for assignments
START_POS_TEXT: Final[Point] = (5, 20) START_POS_TEXT: Final[Point] = (5, 20)
NORMAL_COLOR: Final[ColorBGR] = (255, 0, 255) NORMAL_COLOR: Final[ColorBGR] = (255, 0, 255)
SUCCESS_COLOR: Final[ColorBGR] = (0, 255, 0) SUCCESS_COLOR: Final[ColorBGR] = (0, 255, 0)
FAILURE_COLOR: Final[ColorBGR] = (0, 0, 255) FAILURE_COLOR: Final[ColorBGR] = (0, 0, 255)
FONT_DY: Final[Tuple[int, int]] = (0, cv2.getTextSize("M", FONT, FONT_SCALE, FONT_THICKNESS)[0][1] + 5) FONT_DY: Final[int] = cv2.getTextSize("M", FONT, FONT_SCALE, FONT_THICKNESS)[0][1] + 5
WINDOW_NAME: Final[str] = "Extract OTP Secrets: Capture QR Codes from Camera"
TextPosition = Enum('TextPosition', ['LEFT', 'RIGHT'])
qreader_available = True qreader_available = True
except ImportError: except ImportError:
@ -115,7 +118,7 @@ Otps = List[Otp]
# workaround for PYTHON <= 3.9: OtpUrls = list[OtpUrl] # workaround for PYTHON <= 3.9: OtpUrls = list[OtpUrl]
OtpUrls = List[OtpUrl] OtpUrls = List[OtpUrl]
QRMode = Enum('QRMode', ['QREADER', 'DEEP_QREADER', 'ZBAR', 'CV2', 'WECHAT'], start=0) QRMode = Enum('QRMode', ['QREADER', 'DEEP_QREADER', 'ZBAR', 'CV2', 'CV2_WECHAT'], start=0)
# Constants # Constants
@ -214,6 +217,25 @@ def get_color(new_otps_count: int, otp_url: str) -> ColorBGR:
return NORMAL_COLOR return NORMAL_COLOR
# TODO use cv2 types if available
def cv2_draw_box(img: Any, raw_pts: Any, color: ColorBGR) -> Any:
pts = np.array([raw_pts], np.int32)
pts = pts.reshape((-1, 1, 2))
cv2.polylines(img, [pts], True, color, BOX_THICKNESS)
return pts
# TODO use cv2 types if available
def cv2_print_text(img: Any, text: str, line_number: int, position: TextPosition, color: ColorBGR) -> None:
if position == TextPosition.LEFT:
pos = START_POS_TEXT[0], START_POS_TEXT[1] + line_number * FONT_DY
else:
window_dim = cv2.getWindowImageRect(WINDOW_NAME)
pos = window_dim[2] - cv2.getTextSize(text, FONT, FONT_SCALE, FONT_THICKNESS)[0][0] - 5, START_POS_TEXT[1] + line_number * FONT_DY
cv2.putText(img, text, pos, FONT, FONT_SCALE, color, FONT_THICKNESS, FONT_LINE_STYLE)
def extract_otps_from_camera(args: Args) -> Otps: def extract_otps_from_camera(args: Args) -> Otps:
if verbose: print("Capture QR codes from camera") if verbose: print("Capture QR codes from camera")
otp_urls: OtpUrls = [] otp_urls: OtpUrls = []
@ -222,8 +244,7 @@ def extract_otps_from_camera(args: Args) -> Otps:
qr_mode = QRMode[args.qr] qr_mode = QRMode[args.qr]
cam = cv2.VideoCapture(args.camera) cam = cv2.VideoCapture(args.camera)
window_name = "Extract OTP Secrets: Capture QR Codes from Camera" cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_AUTOSIZE)
cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE)
qreader = QReader() qreader = QReader()
cv2_qr = cv2.QRCodeDetector() cv2_qr = cv2.QRCodeDetector()
@ -236,7 +257,7 @@ def extract_otps_from_camera(args: Args) -> Otps:
break break
try: try:
if qr_mode in [QRMode.QREADER, QRMode.DEEP_QREADER]: if qr_mode in [QRMode.QREADER, QRMode.DEEP_QREADER]:
bbox, found = qreader.detect(img) found, bbox = qreader.detect(img)
if qr_mode == QRMode.DEEP_QREADER: if qr_mode == QRMode.DEEP_QREADER:
otp_url = qreader.detect_and_decode(img, True) otp_url = qreader.detect_and_decode(img, True)
elif qr_mode == QRMode.QREADER: elif qr_mode == QRMode.QREADER:
@ -244,15 +265,13 @@ def extract_otps_from_camera(args: Args) -> Otps:
if otp_url: if otp_url:
new_otps_count = extract_otps_from_otp_url(otp_url, otp_urls, otps, args) new_otps_count = extract_otps_from_otp_url(otp_url, otp_urls, otps, args)
if found: if found:
cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), get_color(new_otps_count, otp_url), RECT_THICKNESS) cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), get_color(new_otps_count, otp_url), BOX_THICKNESS)
elif qr_mode == QRMode.ZBAR: elif qr_mode == QRMode.ZBAR:
for qrcode in zbar.decode(img): for qrcode in zbar.decode(img):
otp_url = qrcode.data.decode('utf-8') otp_url = qrcode.data.decode('utf-8')
pts = numpy.array([qrcode.polygon], numpy.int32)
pts = pts.reshape((-1, 1, 2))
new_otps_count = extract_otps_from_otp_url(otp_url, otp_urls, otps, args) new_otps_count = extract_otps_from_otp_url(otp_url, otp_urls, otps, args)
cv2.polylines(img, [pts], True, get_color(new_otps_count, otp_url), RECT_THICKNESS) cv2_draw_box(img, [qrcode.polygon], get_color(new_otps_count, otp_url))
elif qr_mode in [QRMode.CV2, QRMode.WECHAT]: elif qr_mode in [QRMode.CV2, QRMode.CV2_WECHAT]:
if QRMode.CV2: if QRMode.CV2:
otp_url, raw_pts, _ = cv2_qr.detectAndDecode(img) otp_url, raw_pts, _ = cv2_qr.detectAndDecode(img)
else: else:
@ -260,9 +279,7 @@ def extract_otps_from_camera(args: Args) -> Otps:
if raw_pts is not None: if raw_pts is not None:
if otp_url: if otp_url:
new_otps_count = extract_otps_from_otp_url(otp_url, otp_urls, otps, args) new_otps_count = extract_otps_from_otp_url(otp_url, otp_urls, otps, args)
pts = numpy.array([raw_pts], numpy.int32) cv2_draw_box(img, raw_pts, get_color(new_otps_count, otp_url))
pts = pts.reshape((-1, 1, 2))
cv2.polylines(img, [pts], True, get_color(new_otps_count, otp_url), RECT_THICKNESS)
else: else:
assert False, f"Wrong QReader mode {qr_mode.name}" assert False, f"Wrong QReader mode {qr_mode.name}"
except AssertionError as e: except AssertionError as e:
@ -273,18 +290,13 @@ def extract_otps_from_camera(args: Args) -> Otps:
qr_mode = next_qr_mode(qr_mode) qr_mode = next_qr_mode(qr_mode)
continue continue
cv2.putText(img, f"Mode: {qr_mode.name} (Hit space to change)", START_POS_TEXT, FONT, FONT_SCALE, NORMAL_COLOR, FONT_THICKNESS, FONT_LINE_STYLE) cv2_print_text(img, f"Mode: {qr_mode.name} (Hit space to change)", 0, TextPosition.LEFT, NORMAL_COLOR)
cv2.putText(img, "Hit ESC to quit", tuple(map(add, START_POS_TEXT, FONT_DY)), FONT, FONT_SCALE, NORMAL_COLOR, FONT_THICKNESS, FONT_LINE_STYLE) cv2_print_text(img, "Hit ESC to quit", 1, TextPosition.LEFT, NORMAL_COLOR)
window_dim = cv2.getWindowImageRect(window_name) cv2_print_text(img, f"{len(otp_urls)} QR code{'s'[:len(otp_urls) != 1]} captured", 0, TextPosition.RIGHT, NORMAL_COLOR)
qrcodes_text = f"{len(otp_urls)} QR code{'s'[:len(otp_urls) != 1]} captured" cv2_print_text(img, f"{len(otps)} otp{'s'[:len(otps) != 1]} extracted", 1, TextPosition.RIGHT, NORMAL_COLOR)
pos_qrcodes_text = window_dim[2] - cv2.getTextSize(qrcodes_text, FONT, FONT_SCALE, FONT_THICKNESS)[0][0] - 5, START_POS_TEXT[1]
cv2.putText(img, qrcodes_text, pos_qrcodes_text, FONT, FONT_SCALE, NORMAL_COLOR, FONT_THICKNESS, FONT_LINE_STYLE)
otps_text = f"{len(otps)} otp{'s'[:len(otps) != 1]} extracted" cv2.imshow(WINDOW_NAME, img)
pos_otps_text = window_dim[2] - cv2.getTextSize(otps_text, FONT, FONT_SCALE, FONT_THICKNESS)[0][0] - 5, START_POS_TEXT[1] + FONT_DY[1]
cv2.putText(img, otps_text, pos_otps_text, FONT, FONT_SCALE, NORMAL_COLOR, FONT_THICKNESS, FONT_LINE_STYLE)
cv2.imshow(window_name, img)
key = cv2.waitKey(1) & 0xFF key = cv2.waitKey(1) & 0xFF
if key == 27 or key == ord('q') or key == 13: if key == 27 or key == ord('q') or key == 13:
@ -293,7 +305,7 @@ def extract_otps_from_camera(args: Args) -> Otps:
elif key == 32: elif key == 32:
qr_mode = next_qr_mode(qr_mode) qr_mode = next_qr_mode(qr_mode)
if verbose: print(f"QR reading mode: {qr_mode}") if verbose: print(f"QR reading mode: {qr_mode}")
if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1: if cv2.getWindowProperty(WINDOW_NAME, cv2.WND_PROP_VISIBLE) < 1:
# Window close clicked # Window close clicked
break break
@ -427,9 +439,9 @@ def convert_img_to_otp_url(filename: str, args: Args) -> OtpUrls:
if not stdin: if not stdin:
log_warn("stdin is empty") log_warn("stdin is empty")
try: try:
img_array = numpy.frombuffer(stdin, dtype='uint8') img_array = np.frombuffer(stdin, dtype='uint8')
except TypeError as e: except TypeError as e:
abort(f"Cannot read binary stdin buffer.", e) abort("Cannot read binary stdin buffer.", e)
if not img_array.size: if not img_array.size:
return [] return []
img = cv2.imdecode(img_array, cv2.IMREAD_UNCHANGED) img = cv2.imdecode(img_array, cv2.IMREAD_UNCHANGED)
@ -445,7 +457,7 @@ def convert_img_to_otp_url(filename: str, args: Args) -> OtpUrls:
elif qr_mode == QRMode.CV2: elif qr_mode == QRMode.CV2:
otp_url, _, _ = cv2.QRCodeDetector().detectAndDecode(img) otp_url, _, _ = cv2.QRCodeDetector().detectAndDecode(img)
otp_urls.append(otp_url) otp_urls.append(otp_url)
elif qr_mode == QRMode.WECHAT: elif qr_mode == QRMode.CV2_WECHAT:
otp_url, _ = cv2.wechat_qrcode.WeChatQRCode().detectAndDecode(img) otp_url, _ = cv2.wechat_qrcode.WeChatQRCode().detectAndDecode(img)
otp_urls += list(otp_url) otp_urls += list(otp_url)
elif qr_mode == QRMode.ZBAR: elif qr_mode == QRMode.ZBAR:
@ -672,6 +684,7 @@ def eprint(*args: Any, **kwargs: Any) -> None:
'''Print to stderr.''' '''Print to stderr.'''
print(*args, file=sys.stderr, **kwargs) print(*args, file=sys.stderr, **kwargs)
# workaround for PYTHON <= 3.9 use: BaseException | None # workaround for PYTHON <= 3.9 use: BaseException | None
def abort(msg: str, exception: Optional[BaseException] = None) -> None: def abort(msg: str, exception: Optional[BaseException] = None) -> None:
log_error(msg, exception) log_error(msg, exception)