#! /usr/bin/env python3
# Run without arguments for help message

def help():
    print("Run the tests")
    print("  test.py OMEMO.db")

import sys

# Gajim plugins have garbage debug output, remove it
prev_output = sys.stderr
sys.stderr = open("/dev/null", 'w')

# omemo module is copied from gajim-plugins, assuming axolotl is already installed on system
from omemo.backend.liteaxolotlstore import LiteAxolotlStore, _convert_identity_key as convert
from axolotl.ecc.curve import Curve

# Reestablish STDERR output so we don't eat errors
sys.stderr = prev_output

# Feed me a private key
def sign(key, message):
    return Curve.calculateSignature(key, bytes(message, "utf-8")).hex()

# Feed me a public key
def verify(key, message, sig):
    return Curve.verifySignature(key, bytes(message, "utf-8"), bytes.fromhex(sig))

args = len(sys.argv)
if args == 1:
    help()
    exit(0)

try:
    db = LiteAxolotlStore(sys.argv[1], None)
except:
    print("Failed to load database from " + sys.argv[1])
    exit(1)
keypair = db.getIdentityKeyPair()

# Now let's begin the tests

mymessage = "test message"
sig = sign(keypair.getPrivateKey(), mymessage)

if verify(keypair.getPublicKey().getPublicKey(), mymessage, sig):
    print("OK: Good signature is verified.")
else:
    print("FAIL: Good signature failed to verify!")

# Let's modify the signature
fake_sig = sig.replace('a', 'b').replace('b', 'c')

# Now it should not verify anymore
if verify(keypair.getPublicKey().getPublicKey(), mymessage, fake_sig):
    print("FAIL: Wrong signature is verified!")
else:
    print("OK: Wrong signature is not verified.")

# Let's modify original message
fake_msg = "foobar"
# It should not verify with this fake message and the original signature
if verify(keypair.getPublicKey().getPublicKey(), fake_msg, sig):
    print("FAIL: Wrong message is verified!")
else:
    print("OK: Wrong message is not verified.")