#!/usr/bin/env python3

# Random demo data generator for Lattice ECDSA Attack
# Copyright (C) 2021  Antoine Ferron - BitLogiK
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
#

import argparse
import random
import json

import ecdsa_lib

from fpylll import LLL, BKZ, IntegerMatrix

def reduce_lattice(lattice, block_size=None):
    if block_size is None:
        return LLL.reduction(lattice)
    return BKZ.reduction(
        lattice,
        BKZ.Param(
            block_size=block_size,
            strategies=BKZ.DEFAULT_STRATEGY,
            auto_abort=True,
        ),
    )


def test_result(mat, target_pubkey, curve):
    mod_n = ecdsa_lib.curve_n(curve)
    for row in mat:
        candidate = row[-2] % mod_n
        if candidate > 0:
            cand1 = candidate
            cand2 = mod_n - candidate
            if target_pubkey == ecdsa_lib.privkey_to_pubkey(cand1, curve):
                return cand1
            if target_pubkey == ecdsa_lib.privkey_to_pubkey(cand2, curve):
                return cand2
    return 0


def build_matrix(sigs, curve, num_bits, bits_type, hash_val):
    num_sigs = len(sigs)
    n_order = ecdsa_lib.curve_n(curve)
    curve_card = 2 ** ecdsa_lib.curve_size(curve)
    lattice = IntegerMatrix(num_sigs + 2, num_sigs + 2)
    kbi = 2 ** num_bits
    inv = ecdsa_lib.inverse_mod
    if hash_val is not None:
        hash_i = hash_val
    if bits_type == "LSB":
        for i in range(num_sigs):
            lattice[i, i] = 2 * kbi * n_order
            if hash_val is None:
                hash_i = sigs[i]["hash"]
            lattice[num_sigs, i] = (
                2
                * kbi
                * (
                    inv(kbi, n_order)
                    * (sigs[i]["r"] * inv(sigs[i]["s"], n_order))
                    % n_order
                )
            )
            lattice[num_sigs + 1, i] = (
                2
                * kbi
                * (
                    inv(kbi, n_order)
                    * (sigs[i]["kp"] - hash_i * inv(sigs[i]["s"], n_order))
                    % n_order
                )
                + n_order
            )
    else:
        # MSB
        for i in range(num_sigs):
            lattice[i, i] = 2 * kbi * n_order
            if hash_val is None:
                hash_i = sigs[i]["hash"]
            lattice[num_sigs, i] = (
                2 * kbi * ((sigs[i]["r"] * inv(sigs[i]["s"], n_order)) % n_order)
            )
            lattice[num_sigs + 1, i] = (
                2
                * kbi
                * (
                    sigs[i]["kp"] * (curve_card // kbi)
                    - hash_i * inv(sigs[i]["s"], n_order)
                )
                + n_order
            )
    lattice[num_sigs, num_sigs] = 1
    lattice[num_sigs + 1, num_sigs + 1] = n_order
    return lattice


MINIMUM_BITS = 4
RECOVERY_SEQUENCE = [None, 15, 25, 40, 50, 60]
SIGNATURES_NUMBER_MARGIN = 1.03


def minimum_sigs_required(num_bits, curve_name):
    curve_size = ecdsa_lib.curve_size(curve_name)
    return int(SIGNATURES_NUMBER_MARGIN * 4 / 3 * curve_size / num_bits)


def recover_private_key(
    signatures_data, h_int, pub_key, curve, bits_type, num_bits, loop
):

    # Is known bits > 4 ?
    # Change to 5 for 384 and 8 for 521 ?
    if num_bits < MINIMUM_BITS:
        print(
            "This script requires fixed known bits per signature, "
            f"and at least {MINIMUM_BITS}"
        )
        return False

    # Is there enough signatures ?
    n_sigs = minimum_sigs_required(num_bits, curve)
    if n_sigs > len(signatures_data):
        print("Not enough signatures")
        return False

    loop_var = True
    while loop_var:
        sigs_data = random.sample(signatures_data, n_sigs)

        lattice = build_matrix(sigs_data, curve, num_bits, bits_type, h_int)

        for effort in RECOVERY_SEQUENCE:
            lattice = reduce_lattice(lattice, effort)
            res = test_result(lattice, pub_key, curve)
            if res:
                return res
        loop_var = loop
        if loop:
            print("One more try")

    return 0


def lattice_attack_cli(file_name, loop):
    try:
        with open(file_name, "r") as fdata:
            data = json.load(fdata)
    except FileNotFoundError:
        print(f"Data file '{file_name}' was not found.")
        return
    except IOError:
        print(f"Data file {file_name} can't be accessed.")
        return
    except json.JSONDecodeError:
        print("Data file content is not JSON compatible.")
        return
    message = data.get("message")
    if message:
        hash_int = ecdsa_lib.sha2_int(bytes(message))
    else:
        hash_int = None  # Signal to use a hash per sig, sig data
    curve_string = data["curve"]
    data_type = data["known_type"]
    known_bits = data["known_bits"]
    signatures = data["signatures"]
    q_target = data["public_key"]
    if not ecdsa_lib.check_publickey(q_target, curve_string):
        print(
            f"Public key data invalid, not on the given {curve_string.upper()} curve."
        )
        return
    if loop:
        print("Will shuffle loop until the key found.")
    result = recover_private_key(
        signatures, hash_int, q_target, curve_string, data_type, known_bits, loop
    )
    if result:
        return result
    return 0


def get_p_value():
    return 115792089237316195423570985008687907853269984665640564039457584007908834671663


def add_point(px,py,qx,qy):
    if(px==qx):
        return double_point(px,py)
    p_value=get_p_value()
    almost_inverse=qx-px+p_value
    almost_second=qy-py+p_value
    almost_c=(almost_second*ecdsa_lib.inverse_mod(almost_inverse,p_value))
    c=almost_c%p_value
    rx=((c*c)+(2*p_value)-(px+qx))%p_value
    ry=(c*(px-rx+p_value)+p_value-py)%p_value
    return [rx,ry]


def double_point(px,py):
    p_value=get_p_value()
    triple_px_square=3*px*px
    double_py=2*py
    inverse_double_py=ecdsa_lib.inverse_mod(double_py,p_value)
    c=(triple_px_square*inverse_double_py)%p_value
    c_square=(c*c)%p_value
    minus_2_px=(2*(p_value-px))%p_value
    rx=(c_square+minus_2_px)%p_value
    ry=(c*(px-rx+p_value)+p_value-py)%p_value
    return [rx,ry]   


def multiply_point(px,py,number):
    has_final_point=False
    final_point=[0,0]
    added_point=[px,py]
    while(number!=0):
        if(number%2!=0):
            if not has_final_point:
                final_point=added_point
                has_final_point=True
            else: final_point=add_point(final_point[0],final_point[1],added_point[0],added_point[1])
        number=number//2
        added_point=double_point(added_point[0],added_point[1])
    return final_point


def generates_signatures(curve):
    puzzle120=int("00000000000000000000000000000000007fffffffffffffffffffffffffffff",16)
    known_bits=16
    n_value=ecdsa_lib.curve_n(curve)
    Q_point=[93499419120076195219278579763555015417347613618260420189054155605804414805552,19494200143356336257404688340550956357466777176798681646526975620299854296492]
    sigs = []
    for index in range(100):
        binaryIndex=(index*2).to_bytes(32,'big')
        binaryIndex2=(index*2+1).to_bytes(32,'big')
        hashedIndex=ecdsa_lib.sha2_int(binaryIndex)
        hashedIndex2=ecdsa_lib.sha2_int(binaryIndex2)
        #print(binaryIndex.hex(),hex(hashedIndex))
        #print(binaryIndex2.hex(),hex(hashedIndex2))
        z_value_with_r=(hashedIndex%puzzle120) #random.randrange(puzzle120)
        added_point=ecdsa_lib.privkey_to_pubkey(z_value_with_r,curve)
        final_point=add_point(Q_point[0],Q_point[1],added_point[0],added_point[1])
        r_value_with_s=(hashedIndex2%puzzle120) #random.randrange(puzzle120)
        print(hex(z_value_with_r),hex(r_value_with_s))
        R_point=multiply_point(final_point[0],final_point[1],r_value_with_s)
        r_value=R_point[0]
        s_value_with_r=ecdsa_lib.inverse_mod(r_value_with_s,n_value)
        s_value=(s_value_with_r*r_value)%n_value
        z_value=(z_value_with_r*r_value)%n_value
        sigs.append(
            {
                "r": r_value,
                "s": s_value,
                "kp": 0,
                "hash": z_value
            }
        )
    ret = {
        "curve": curve.upper(),
        "public_key": Q_point,
        "known_type": "MSB",
        "known_bits": known_bits,
        "signatures": sigs,
    }
    return ret


if __name__ == "__main__":
    curve="secp256k1"
    n_value=ecdsa_lib.curve_n(curve)
    sigs_data = generates_signatures(curve)
    with open("data.json", "w") as fout:
        json.dump(sigs_data, fout)   
    '''
    d_value=lattice_attack_cli("data.json", False)
    if(d_value==0):
        print("failed")
    else:
        r_value=sigs_data["signatures"][0]["r"]
        s_value=sigs_data["signatures"][0]["s"]
        z_value=sigs_data["signatures"][0]["hash"]
        rd_value=(r_value*d_value)%n_value
        rd_plus_z_value=(rd_value+z_value)%n_value
        inverted_s_value=ecdsa_lib.inverse_mod(s_value,n_value)
        k_value=(rd_plus_z_value*inverted_s_value)%n_value
        print("privkey:",hex(k_value))
    '''