""" Class defintion for simple hastad broadcast exploit """

from RSAExploits import util
from RSAExploits.exploits.exploit import Exploit

class Hastad(Exploit):
    """ Class providing a run interface to hastad broadcast exploit"""
    
    def run(self, rsadata_list, info_dict = None):
        """ Attempts to recover plaintext using Hastad's broadcast attack 

            This attack works when we have a list of RSA_Data objects such
            that the length of the list is greater than or equal to e, 
            unique and coprime moduli are used for each encryption, and the
            same plaintext message is encrypted to generate all ciphertexts.
    
            Args:
                rsadata_list: A list of RSA_Data objects on which to attempt
                              Hastad's exploit
                info_dict: Not currently used

            Assumptions:
                None of the RSA_Data objects in rsadata_list have the same 
                public exponent e, same id number, and same modulus N. This 
                should be prevented by calling the remove_duplicates() function 
                in RSA_Data.py

            Side Effects: 
                If a message is recovered, the corresponding RSA_Data objects
                will be updated with this information

            Return: 
                True if at least one message was recovered
        """
        
        print ("Hastad: Running Attack...")
        success = False     
        e_id_dict = self.group_by_e_and_id(rsadata_list)
        for group in e_id_dict:
            msg = self.hastad_broadcast_exploit(e_id_dict[group])
            if msg != None:
                success = True 
                for rsadata in e_id_dict[group]:
                    rsadata.set_m(msg)              
        if success:
            print("Hastad: Success, message found.")                            
        else:           
            print("Hastad: Failure, message not found.")
        return success
            
    @staticmethod
    def group_by_e_and_id(rsadata_list):
        """ Group the RSA_Data objects by public exponent and id """
 
        e_id_dict = {}
    
        for rsadata in rsadata_list:
        
            # Only consider entries with an ID number
            if rsadata.get_id() == None:
                continue

            # Only consider entries with a ciphertext
            if rsadata.get_c() == None:
                continue

            # If the (e, idnum) tuple already exists in the dictionary, just
            # append the new entry to the already existing list
            if (rsadata.get_e(), rsadata.get_id()) in e_id_dict:
                e_id_dict[(rsadata.get_e(), rsadata.get_id())].append(rsadata)
		
            # Otherwise, create a new list for the new tuple
            else:
                e_id_dict[(rsadata.get_e(), rsadata.get_id())] = [rsadata]

        return e_id_dict
    

    @staticmethod
    def hastad_broadcast_exploit(rsadata_list):
        """ Recover the plaintext message using chinese remainder theorem """   
       
        e = rsadata_list[0].get_e() 
        if len(rsadata_list) < e:
            return None

        ns = []
        cs = []

        for index in range(e):
            ns.append(rsadata_list[index].get_n())
            cs.append(rsadata_list[index].get_c())

        s = util.crt(ns, cs)
        pt = util.int_nthroot(s, e)
        if pt is not None:
            return pt
        else:
            return None

            
