petya_decryptor.py

#!/usr/bin/python2

"""
Petya Decryptor

Petya encrypts the MFT of the NTFS partition to prevent the user from accessing files.
It uses a flawed implementation of the Salsa20 stream cipher to perform the encryption.
Since Petya runs in 16-bit real mode, it cannot simply put 32-bit values into 16-bit registers.
To force the copied algorithm to work, the developers simply changed a few constant values of 32 down to 16.
This slight change reduces the keyspace by making a number of values in the keystream predictable.
The 8-byte IV used is saved 33 bytes into sector 54 of the disk.
A 16-byte random key, which is the value we need to solve for, is key-expanded to 32-bytes.
When Petya is provided with a decryption key, it first attempts to decrypt the contents of sector 55.
If the decryption is successful, all the decrypted data will be equal to 0x37.
We create a 16-byte BitVector to represent the decryption key and add our constraints.

References:
- https://0xec.blogspot.com/2016/04/reversing-petya-ransomware-with.html
- https://blog.checkpoint.com/2016/04/11/decrypting-the-petya-ransomware/
- https://en.wikipedia.org/wiki/Salsa20
"""

from z3 import *

BLOCK_SIZE = 512


class Drive(object):
    """
    Layout of the drive's sector map:
    0000h-0200h: New MBR, partition entries stay unchanged, the bootstrap code is changed to load the new bootloader code in "main core"
    0200h-4400h: Data XOR'ed with 0x37
    4400h-6c00h: Main core
    6c00h-6e00h: "Onion Sector", contains the encryption key, IV, and other information that identifys the victim host
    6e00h-7000h: Buffer used for verification of decryption key
    7000h-7200h: Original MBR XOR'ed with 0x37
    """

    def __init__(self, filename):
        self.filename = filename

    def read_sector(self, sector, key=None):
        start = BLOCK_SIZE * sector

        with open(self.filename, 'rb') as f:
            f.seek(sector * BLOCK_SIZE)
            data = bytearray(f.read(BLOCK_SIZE))

        if key is not None:
            for i in range(len(data)):
                data[i] ^= key

        return data

    def new_mbr(self):
        return self.read_sector(0)

    def orig_mbr(self):
        return self.read_sector(56, 0x37)

    def onion_sector(self):
        return self.read_sector(54)

    def is_encrypted(self):
        data = self.read_sector(54)
        return ord(data[0]) == 1

    def pubkey(self):
        data = self.read_sector(54)
        return data[0xA9: 0xA9 + 0x5A]

    def get_iv(self):
        data = self.read_sector(54)
        return data[33: 33 + 8]


# Rotate-left operation for Salsa20 used input values of 32-bits
# The Petya implementation operates in 16-bit real mode,
# which means the value is 16, not 32, yet we still subtract it from 32
def rotl(value, shift):
    return (value << shift) | LShR(value, 32 - shift)


def quarterround(y0, y1, y2, y3):
    y1 = y1 ^ rotl(y0 + y3, 7)
    y2 = y2 ^ rotl(y1 + y0, 9)
    y3 = y3 ^ rotl(y2 + y1, 13)
    y0 = y0 ^ rotl(y3 + y2, 18)
    return y0, y1, y2, y3


def rowround(y):
    y[0], y[1], y[2], y[3] = quarterround(y[0], y[1], y[2], y[3])
    y[5], y[6], y[7], y[4] = quarterround(y[5], y[6], y[7], y[4])
    y[10], y[11], y[8], y[9] = quarterround(y[10], y[11], y[8], y[9])
    y[15], y[12], y[13], y[14] = quarterround(y[15], y[12], y[13], y[14])
    return y


def columnround(x):
    x[0], x[4], x[8], x[12] = quarterround(x[0], x[4], x[8], x[12])
    x[5], x[9], x[13], x[1] = quarterround(x[5], x[9], x[13], x[1])
    x[10], x[14], x[2], x[6] = quarterround(x[10], x[14], x[2], x[6])
    x[15], x[3], x[7], x[11] = quarterround(x[15], x[3], x[7], x[11])
    return x


def doubleround(x):
    x = columnround(x)
    x = rowround(x)
    return x


def little_endian(b0, b1):
    """Swaps the given bytes to little endian order"""
    return b1, b0


def reverse_little_endian(b, w, i):
    b[i + 0] = Extract(7, 0, w)
    b[i + 1] = Extract(15, 8, w)
    b[i + 2] = BitVecVal(0, 8)
    b[i + 3] = BitVecVal(0, 8)
    return b


def hash(seq):
    # Petya's developers incorrectly reduced the size of these variables
    # from 32 to 16 bits in order to function in 16-bit real mode
    x = [BitVec('x_%d' % i, 16) for i in range(16)]
    z = [BitVec('z_%d' % i, 16) for i in range(16)]

    for i in range(16):
        # Petya reads 2-bytes here but still increments by 4 bytes
        # This flaw forces it to generate a 512-bit keystream with
        # 256 predictable bits, thus reducing the keyspace
        b0, b1 = little_endian(seq[4 * i], seq[4 * i + 1])
        x[i] = z[i] = Concat(BitVecVal(0, 8), b1) + Concat(b0, BitVecVal(0, 8))

    for i in range(10):
        z = doubleround(z)

    for i in range(16):
        z[i] += x[i]
        seq = reverse_little_endian(seq, z[i], 4 * i)

    return seq


def expand32(k, n, keystream):
    # The constants specified by the Salsa20 specification, 'sigma' "expand 32-byte k"
    o = [
        ['e', 'x', 'p', 'a'],
        ['n', 'd', ' ', '3'],
        ['2', '-', 'b', 'y'],
        ['t', 'e', ' ', 'k']
    ]

    # Copy all of 'sigma' into the correct spots in the keystream block
    for i in range(0, 64, 20):
        for j in range(4):
            keystream[i + j] = BitVecVal(ord(o[i / 20][j]), 8)

    # Copy the key and the IV into the keystream block
    for i in range(16):
        keystream[4 + i] = k[i]
        keystream[44 + i] = k[i + 16]
        keystream[24 + i] = n[i]

    keystream = hash(keystream)
    return k, n, keystream


if __name__ == "__main__":
    drive = Drive("/home/binaries/files/ransom.img")

    # The eight byte IV is located in the "Onion Sector" (sector 54)
    # IV is expanded to 16 bytes by appending eight 0's
    iv = drive.get_iv()
    iv.extend(bytearray([0] * 8))
    print("The IV recovered from sector 54 is: " + str(iv).encode("hex"))

    # Read the verification bytes from sector 55 and xor them with 0x37
    # Expanding the key must produce this keystream, i.e our constraint is that
    # the keystream generated by the correct encryption key equals this keystream
    target_keystream = drive.read_sector(55, 0x37)[:64]
    print("The target keystream recovered from sector 55 is: " +
          str(target_keystream).encode("hex"))

    key = []
    bvlist = []
    n = [BitVecVal(i, 8) for i in iv]  # Convert each byte in the IV to BitVecs
    keystream = [0] * 64

    # Key Expansion
    # The original private key is a randomly generated 16-byte value
    # which is then expanded to 32-bytes using the following key expansion.
    for i in range(16):
        bv = BitVec('k_%d' % i, 8)
        bvlist.append(bv)
        key.append(bv + 0x7A)
        key.append(bv * 2)

    key, n, keystream = expand32(key, n, keystream)
    s = Solver()

    # Constrain every byte in the keystream to be equal to the known target keystream
    for i in range(len(target_keystream)):
        s.add(keystream[i] == target_keystream[i])

    # Solve for the encryption key
    if s.check() == sat:
        m = s.model()
        keystring = reduce(
            lambda acc, x: acc + '1' if m[x] == None else acc + chr(m[x].as_long()), bvlist, '')
        print("Decryption Key: " + str(keystring))