txlyre 8 months ago
commit
b2d6379531
3 changed files with 274 additions and 0 deletions
  1. 3 0
      .gitignore
  2. 267 0
      main.py
  3. 4 0
      requirements.txt

+ 3 - 0
.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+venv
+.secret

+ 267 - 0
main.py

@@ -0,0 +1,267 @@
+import os
+import os.path
+import time
+import argparse
+import tempfile
+import traceback
+
+import cv2
+import requests
+from loguru import logger
+from Crypto.Hash import BLAKE2b
+
+try:
+    import rdrand
+except ImportError:
+    logger.warning("RdRand is not available.")
+
+
+class TemporaryFile:
+    def __init__(self, name, io, delete):
+        self.name = name
+        self.__io = io
+        self.__delete = delete
+
+    def __getattr__(self, k):
+        return getattr(self.__io, k)
+
+    def __del__(self):
+        if self.__delete:
+            try:
+                os.unlink(self.name)
+            except FileNotFoundError:
+                pass
+
+
+def NamedTemporaryFile(
+    mode="w+b", bufsize=-1, suffix="", prefix="tmp", dir=None, delete=True
+):
+    if not dir:
+        dir = tempfile.gettempdir()
+
+    name = os.path.join(dir, prefix + os.urandom(32).hex() + suffix)
+
+    if mode is None:
+        return TemporaryFile(name, None, delete)
+
+    fh = open(name, "w+b", bufsize)
+    if mode != "w+b":
+        fh.close()
+        fh = open(name, mode)
+
+    return TemporaryFile(name, fh, delete)
+
+
+def chunks(lst, n):
+    for i in range(0, len(lst), n):
+        yield lst[i : i + n]
+
+
+def run_check(cmd):
+    logger.info(f"Executing '{cmd}'.")
+
+    if os.system(cmd) != 0:
+        raise ValueError("Exit code != 0.")
+
+
+def extract_image(path):
+    logger.info(f"Extract image '{path}'.")
+
+    im = cv2.imread(path)
+
+    data = []
+    rows, cols, _ = im.shape
+    for i in range(rows):
+        for j in range(cols):
+            r, g, b = im[i, j]
+
+            data.append(((r << 16) + (g << 8) + b) & 255)
+
+    return bytes(data)
+
+
+def extract_wav(path):
+    data = []
+
+    with open(path, "rb") as f:
+        for sample in chunks(f.read()[44:], 2):
+            data.append(sample[0])
+
+    return bytes(data)
+
+
+def extract_lsbs(data):
+    logger.info("Extract LSBs.")
+
+    buffer = []
+
+    if len(data) % 2 != 0:
+        data = data[:-1]
+
+    for chunk in chunks(data, 2):
+        tmp_byte = 0
+        for byte in chunk:
+            for n in range(4):
+                tmp_byte = (tmp_byte << 1) | ((byte >> n) & 1)
+
+        buffer.append(tmp_byte & 255)
+
+    return bytes(buffer)
+
+
+def whiten(data):
+    logger.info("Whitening.")
+
+    buffer = b""
+    for chunk in chunks(data, 128):
+        buffer += BLAKE2b.new(data=chunk, digest_bits=256).digest()
+
+    return buffer
+
+
+def read_video(source, duration=60):
+    tmpf = NamedTemporaryFile(suffix=".mkv", mode=None)
+
+    run_check(
+        f"ffmpeg -hide_banner -loglevel error -y -i {source} -t {duration} -acodec copy -vcodec copy {tmpf.name}"
+    )
+
+    with tempfile.TemporaryDirectory() as tmpd:
+        run_check(
+            f"ffmpeg -hide_banner -loglevel error -y -i {tmpf.name} -vf mpdecimate -r 1/1 {tmpd.name}/%d.bmp"
+        )
+
+        data = b""
+        for filename in os.listdir(tmpd.name):
+            data += extract_image(os.path.join(tmpd.name, filename))
+
+        return data
+
+
+def read_audio(source, duration=60):
+    tmpf = NamedTemporaryFile(suffix=".wav", mode=None)
+
+    run_check(
+        f"ffmpeg -f alsa -i {source} -t {duration} -ar 48000 -f s16le -acodec pcm_s16le {tmpf.name}"
+    )
+
+    return extract_wav(tmpf.name)
+
+
+def read_audio_video(source, duration=60):
+    tmpf = NamedTemporaryFile(suffix=".mkv", mode=None)
+
+    run_check(
+        f"ffmpeg -hide_banner -loglevel error -y -i {source} -t {duration} -acodec copy -vcodec copy {tmpf.name}"
+    )
+
+    with tempfile.TemporaryDirectory() as tmpd:
+        run_check(
+            f"ffmpeg -hide_banner -loglevel error -y -i {tmpf.name} -vf mpdecimate -r 1/1 {tmpd}/%d.bmp"
+        )
+
+        data_a = b""
+        for filename in sorted(
+            os.listdir(tmpd), key=lambda filename: int(filename.split(".")[0])
+        ):
+            data_a += extract_image(os.path.join(tmpd, filename))
+
+        tmpf2 = NamedTemporaryFile(suffix=".wav", mode=None)
+
+        run_check(
+            f"ffmpeg -hide_banner -loglevel error -y -i {tmpf.name} -vn -ar 48000 -f s16le -acodec pcm_s16le {tmpf2.name}"
+        )
+
+        data_b = extract_wav(tmpf2.name)
+
+        return bytes(a ^ b for a, b in zip(data_a, data_b))
+
+
+def read_rdseed(source, amount=16):
+    data = rdrand.rdseed_get_bytes(amount)
+    if len(data) != amount or data.count(0) == amount:
+        raise ValueError("bad data")
+
+    return data
+
+
+def sample(source, source_type, multiplier=1):
+    match source_type:
+        case "video":
+            sampler = read_video
+            multiplier *= 60
+
+        case "audio":
+            sampler = read_audio
+            multiplier *= 60
+
+        case "video+audio":
+            sampler = read_audio_video
+            multiplier *= 60
+
+        case "rdseed":
+            sampler = read_rdseed
+
+        case _:
+            raise ValueError(source_type)
+
+    multiplier = int(multiplier)
+    if multiplier < 1:
+        raise ValueError(multiplier)
+
+    logger.info("Sampling...")
+
+    data = sampler(source, multiplier)
+
+    logger.info(f"Sample ready: {len(data)}b.")
+
+    if source_type != "rdseed":
+        data = extract_lsbs(data)
+        data = whiten(data)
+
+    return data
+
+
+def push(pool_url, data, secret):
+    resp = requests.post(
+        f"{pool_url}/api/pool", data=data[: 1024**2], headers={"X-Secret": secret}
+    )
+
+    (logger.success if resp.status_code == 200 else logger.error)(
+        f"{resp.status_code}: {resp.text}"
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", type=str, required=True)
+    parser.add_argument("--source-type", type=str, default="video+audio")
+    parser.add_argument("--multiplier", type=float, default=1)
+    parser.add_argument("--secret-file", type=str, default="./.secret")
+    parser.add_argument("--cooldown", type=int, default=0)
+    parser.add_argument("--pool-url", type=str, default="https://trng.iike.ru")
+
+    args = parser.parse_args()
+
+    with open(args.secret_file, "r") as f:
+        lines = f.read().strip().split("\n")
+        ident = lines[0].strip()
+        secret = lines[1].strip()
+
+        secret = f"{ident} {secret}"
+
+    while True:
+        try:
+            data = sample(args.source, args.source_type, args.multiplier)
+
+            push(args.pool_url, data, secret)
+        except KeyboardInterrupt:
+            logger.info("Interrupted by user.")
+
+            break
+        except Exception as e:
+            traceback.print_exc()
+
+            logger.error(e)
+
+        time.sleep(args.cooldown)

+ 4 - 0
requirements.txt

@@ -0,0 +1,4 @@
+loguru
+requests
+pycryptodome
+opencv-python