Old script use to extract shared prime from a number of public RSA keys (i.e. generated by malware)

extractrsa.py

#!/usr/bin/env python3
#
# author: notdodo
#
import os
import sqlite3
import itertools
from multiprocessing import Pool

try:
    import gmpy2
    from Crypto.PublicKey import RSA
    from factordb.factordb import FactorDB
    from colored import fg, stylize
except:
    print("pip install factordb-pybcli gmpy2 colored pycrypto")
    import sys

    sys.exit(-1)


##############################################################################
#
# SQLite3 support functions
#
def db_connect(base_path, db_name):
    db_name = os.path.abspath(os.path.join(base_path, db_name))
    print("DB Path: {}".format(db_name))
    return sqlite3.connect(db_name)


def db_close(db_conn):
    db_conn.commit()
    db_conn.close()


##############################################################################


#
# Save a list of RSA key (PEM) from a directory to a SQLite3 DB
# TODO: make it parallel
#
def extract_save(base_path, db_name):
    db = db_connect(base_path, db_name)
    cursor = db.cursor()

    cursor.execute(
        """CREATE TABLE IF NOT EXISTS keys_infos
        (filename TEXT PRIMARY KEY, modulus TEXT, exponent TEXT)"""
    )

    file_number = len(os.listdir(base_path))
    count = 0

    # Extract informations from the key
    def pem_info(pub_key):
        key = RSA.importKey(pub_key)
        return key.n, key.e

    for key in os.listdir(base_path):
        file_key = os.path.join(base_path, key)

        try:
            modulus, exponent = pem_info(open(file_key).read())

            cursor.execute(
                "INSERT OR IGNORE INTO keys_infos VALUES ('{}', '{}', '{}')".format(
                    key, modulus, exponent
                )
            )

            # Test if the modulus is a known factor (TODO: add/find other services)
            factors = FactorDB(modulus).get_factor_list()
            if len(factors) > 1:
                print(
                    """
                    [************************************************************]\n
                    {}\n
                    [************************************************************]
                    """.format(
                        factors
                    )
                )
        except Exception as e:
            print(type(e), e, file_key)
            continue
        else:
            # Print progress
            count += 1
            print(
                "[-] Completion: {}/{}; Current: {}".format(
                    stylize(count, fg("red")),
                    stylize(file_number, fg("blue")),
                    stylize(key, fg("green")),
                ),
                end="\r",
            )

    print("\n[*] All keys saved!")
    # Save (commit) the changes
    db_close(db)


def gcd(pair):
    # gdb(a, b) = [mpz(), a, b]
    return [gmpy2.gcd(int(pair[0]), int(pair[1])), pair[0], pair[1]]


def run_gcd(base_path, db_name):
    modulus = []
    modulus_index = 2
    group_slice = 1000
    db = db_connect(base_path, db_name)
    cursor = db.cursor()

    # TODO: try generator to improve performance and reduce memory occupation
    for row in cursor.execute("SELECT rowid, * FROM keys_infos"):
        modulus.append(row[modulus_index])

    # split iterable object in slices of n elements
    def grouper(n, iterable, fill=(1, 1)):
        args = [iter(iterable)] * n
        return itertools.zip_longest(fillvalue=fill, *args)

    # Generate all possibile combinations of permutions of 2 elements
    comb = itertools.combinations(modulus, 2)
    # Will contain shared keys
    shared = set()
    with Pool() as p:
        for i in grouper(group_slice, comb):
            result = []
            result = [x for x in p.map(gcd, i) if x[0] != 1]

            # Clean result and find only real shared prime numbers
            # TODO: improve this step: IS UGLY!
            for r in result:
                [
                    shared.add(fn[0])
                    for fn in cursor.execute(
                        """SELECT filename FROM keys_infos WHERE
                            modulus in ("{}", "{}")""".format(
                            r[1], r[2]
                        )
                    )
                ]

    if len(shared):
        print(shared)
        parent_dir = os.path.abspath(os.path.join(base_path, os.pardir))
        with open(
            os.path.abspath(os.path.join(parent_dir, "output_shared.txt")), "w+"
        ) as save:
            save.writelines(x for x in shared)
    else:
        print(stylize("[!!] No shared prime numbers found", fg("red")))
    db_close(db)


if __name__ == "__main__":
    keys_path = "/home/dodo/Downloads/pk/"
    db_name = "pb_keys.db"
    extract_save(keys_path, db_name)
    run_gcd(keys_path, db_name)