import base64
import binascii
import os
import struct

'''
required: sudo pip install pycrypto
'''
from Crypto.Util.number import bytes_to_long, long_to_bytes
from Crypto.Hash import SHA256
from Crypto import Random
from Crypto.Cipher import AES
from Crypto.Util import Counter

'''
required: sudo apt-get install libgmp3-dev && sudo pip install gmpy
'''
from gmpy import invert

'''
WannaCry decrypter from the original script
'''

MAGIC = 0xBEE4FAC3

class Decrypter:

  def _decrypt_header(self, filename, input_file, key):
    iv = input_file.read(AES.block_size)

    if len(iv) < AES.block_size:
      raise Exception("Incorrect len of iv {0}".format(filename))

    ctr = Counter.new(AES.block_size * 8, initial_value=long(binascii.hexlify(iv), 16))
    aes = AES.new(key, mode=AES.MODE_CTR, counter=ctr)

    data = input_file.read(AES.block_size)

    if len(data) < AES.block_size:
      raise Exception("Incorrect len of data {0}".format(filename))

    data = aes.decrypt(data)
    if long(binascii.hexlify(data[:4]), 16) != MAGIC:
      raise Exception("No Magic in file {0}".format(filename))

    lsize = struct.calcsize('<I')
    l = struct.unpack('<I', data[4:4 + lsize])[0]
    filename = str(data[4 + lsize:])
    filename = filename[:l]
    while len(filename) < l:
      data = input_file.read(AES.block_size)
      filename += aes.decrypt(data)
      filename = filename[:l]

    return (aes, filename)

  def _decrypt_file(self, path, filename, key):
    with open(os.path.join(path, filename), 'rb') as input_file:

      aes, output_filename = self._decrypt_header(filename, input_file, key)

      with open(os.path.join(path, output_filename), 'wb') as output_file:
        while True:
          data = input_file.read(AES.block_size)
          if len(data) == 0:
            break
          output_file.write(aes.decrypt(data))
    
    print "Decrypted {0}".format(output_filename)
    os.remove(os.path.join(path, filename))

  def process(self, path, digest_key_list):
    lookup = dict(digest_key_list)
    lookup.update(dict((digest[:32], key) for digest, key in digest_key_list))
    for dirpath, _, filenames in os.walk(path):
      for filename in filenames:
        try:
          key = lookup[filename[1:-7]] 
          self._decrypt_file(dirpath, filename, base64.b64decode(key))
        except KeyError:
          continue

'''
Franklin-Reiter attack function from https://github.com/sonickun/cryptools
'''
def franklin_reiter_related_message_attack(e, n, c1, c2, a, b):
    assert e == 3 and b != 0
    frac = b * (c2 + 2*pow(a,3)*c1 - pow(b,3))
    denom = a * (c2 - pow(a,3)*c1 + 2*pow(b,3))
    m = (frac * invert(denom, n)) % n
    return m

'''
The large modulus number used in the public key of RSA
'''
N = 20734601011924435778622626352029859267407020070459948494552124371885551349803628688810359950660469611598599270653643020199849673645957114359068856671579941583787010742883379752422179790775854722285250574171017314543776132627340992715862463976226157198802724737916167351743031092170694732212079270288260116920347005873441934844331366977833414936053188083278692048655806492237086833973593048029971953858708564756986112746264342160128533820139843707887554443829740749341457682129448431307376277626946193349841459871135143174332865644871124709386593795444594176571398306854978531773059317408297544411480754278445963594097L

'''
The value of the first message sent over UDP converted to long. This is RSA encoded key that is hashed and used as AES encryption key for the file.
'''
C1 = 18640377020667858103936371537153243192953126452150315104918681393098732451816631827310383718851904186643913272781881354929922643123623275580346335489833686775415079675700672132909515165780733643494087663051855493035731830436615169486595010689952614371506059583394789987850989306070748959960531074294143836487310025945815011804630762622821872040212104540034988158810508344403837815896077485172183090588353613848175809464341574915998943170306472309204749017958883095927877004682225973914063170861315797490959700392724190395087888784398603012676378506452137471691831326752109908137454913977632214911809914621268196324932L

'''
The value of the second message sent over UDP converted to long. This is RSA encoded key that is hashed and used as AES encryption key for the file.
'''
C2 = 19166548133369310411210186782125407314073499566019561339865399475306561430640164667166265124865084621444683065835026501516193277506399509979696933843415437981139169186519334598147696048161051824045783469135775741889098388879776083695365317229612113818353137564369271547860302882695872209374209241382107456509463714402392293395942565513509318469331704092298104619169169882260041750432661620840239458808449502048112547192844952759345762626612001729128712490650421045400331691540704482547912240138082830303138038200602795851471358574913481412976664622346204488888284858501352158262062458201349865909650163532469195278595L

'''
Key for the first message
'''
M = franklin_reiter_related_message_attack(3, N, C1, C2, 1, 1)

decrypter = Decrypter()

keys = []
digests = []

'''
compute keys for files and compute their digests to match a key with the file it's been used on
'''
for i in range(0,22):
    key = base64.b64encode(SHA256.new(long_to_bytes(M + i)).digest())
    digest = binascii.hexlify(SHA256.new(base64.b64encode(long_to_bytes(pow(M + i, 3, N)))).digest())
    keys.append(key)
    digests.append(digest)

decrypter.process(os.getcwd(), zip(digests, keys))

'''
Clean up
'''
print "Removing README files with the ransom..."
for d in digests:
  os.remove(os.path.join(os.getcwd(), d[:32] + ".README.txt"))