Mersenne Twister - Python Module

다양한 것들을 시도하던 중에, 예전에 풀다가 포기했던 이론들을 다시 차례차례 접근해야겠다는 생각이 들어서 예전에 시도했던 것들을 하나하나 다시 상기시키는 중이다.

그 중에서도 예전에 Dreamhack 사이트에서 접해보았던 Unbreakable 이라는 문제를 보고, Python random module 의 원리인 Mersenne Twister 에 대해서 공부해보고자 이번 포스팅을 진행하게 되었다. (무려 Lv. 8 짜리 문제이니 긴장하자)

아이러니하게도 또 RBTree 님 문제;;

🖲️ Code Analysis

#!/usr/bin/env python3

from hashlib import sha256
import os
import random
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from flag import iflags, final_flag

TOTAL_BITS = 624 * 37
def challenge(num_bit, previous, current):
    random.seed(os.urandom(16))
    num_task = (TOTAL_BITS + num_bit - 1) // num_bit

    output = ""
    for _ in range(num_task):
        output += format(random.getrandbits(num_bit), f"0{num_bit}b")
    
    print(output)

    key = previous
    for _ in range(100):
        key += format(random.getrandbits(num_bit), f"0{num_bit}b").encode()
        key = sha256(key).digest()
    
    cipher = AES.new(key, AES.MODE_CBC, iv=b'iluvredblacktree')
    print(cipher.encrypt(pad(current, 16)).hex())

challenge(32, b'iloveredblacktree', iflags[0])
challenge(16, iflags[0], iflags[1])
challenge(8, iflags[1], iflags[2])
challenge(4, iflags[2], iflags[3])
challenge(2, iflags[3], iflags[4])
challenge(1, iflags[4], final_flag)

코드를 해석해보면, 상당히 간단 (하지만 어려운) 코드를 통해서 flag 를 얻어낼 수 있다. 

1. num_bit 값을 통해서 num_task 의 값을 계산해내는데, 624 * 37 이라는 값을 사용하여 최소 624 개의 task 가 만들어지게 된다. 

2. 그 후에, random 한 seed 로 만들어낸 getrandbits 를 task 의 수만큼 계산하고 출력하게 된다. 

3. 그 후에 계속해서 random bits 들을 만들어내게 되는데, 일종의 Merkle Damgard 의 형태로 key 를 만ㄷ르어내게 된다. 이 key 를 통해서 flag 를 암호화하게 된다. 

문제를 사실상 Python 의 Random module 의 취약점을 통해서, 앞선 random 값들로 그 후의 random 값들을 예측해낼 수 있겠냐는 질문이 되겠다. 그렇다면 Python 의 Random module 은 어떤 알고리즘을 따를까??

💡 Main Idea

우리는 여기서 Mersenne Twister 에 대한 공부를 할 필요가 있다. 이는 python module 에서 랜덤값들을 만들어내는 알고리즘인데, 이 알고리즘에는 엄청난 취약점이 존재한다.

 

먼저 Mersenne Twister 의 알고리즘을 살펴보자. 

😎 Mersenne Twister Algorithm

/* Period parameters -- These are all magic.  Don't change. */
#define N 624
#define M 397
#define MATRIX_A 0x9908b0dfU    /* constant vector a */
#define UPPER_MASK 0x80000000U  /* most significant w-r bits */
#define LOWER_MASK 0x7fffffffU  /* least significant r bits */

static uint32_t
genrand_uint32(RandomObject *self)
{
    uint32_t y;
    static const uint32_t mag01[2] = {0x0U, MATRIX_A};
    /* mag01[x] = x * MATRIX_A  for x=0,1 */
    uint32_t *mt;

    mt = self->state;
    if (self->index >= N) { /* generate N words at one time */
        int kk;

        for (kk=0;kk<N-M;kk++) {
            y = (mt[kk]&UPPER_MASK)|(mt[kk+1]&LOWER_MASK);
            mt[kk] = mt[kk+M] ^ (y >> 1) ^ mag01[y & 0x1U];
        }
        for (;kk<N-1;kk++) {
            y = (mt[kk]&UPPER_MASK)|(mt[kk+1]&LOWER_MASK);
            mt[kk] = mt[kk+(M-N)] ^ (y >> 1) ^ mag01[y & 0x1U];
        }
        y = (mt[N-1]&UPPER_MASK)|(mt[0]&LOWER_MASK);
        mt[N-1] = mt[M-1] ^ (y >> 1) ^ mag01[y & 0x1U];

        self->index = 0;
    }

    y = mt[self->index++];
    y ^= (y >> 11);
    y ^= (y << 7) & 0x9d2c5680U;
    y ^= (y << 15) & 0xefc60000U;
    y ^= (y >> 18);
    return y;
}

코드가 살짝 복잡해 보이지만, 크게 2가지 방식으로 코드가 진행되는 것을 볼 수 있다.

1. 한번 함수를 호출할 때 마다, mt [self-> index] 를 참조하고,, self -> index 의 값을 1 증가시킨다.
2. 만약 self -> index 가 N, 즉 624 를 넘어섰다면 kk 를 0부터 N-1 까지 진행하는데, 그 과정은 아래와 같다.

y = (mt[kk] & UPPER_MASK) | (mt[kk+1] & LOWER_MASK);

mt[kk] = mt[(kk + M) % N] ^ (y >> 1) ^ mag01[y & 0x1U];

위에 적힌 코드는 복잡해 보이지만, 잘 생각해보면 모두 N 으로 나눈 나머지를 의미한다는 것을 알 수 있을 것이다.

 

결국, 우리가 주요하게 보아야 할 부분은 아래의 암호화 부분이다. 이 부분은 temper 이라고 부르는데, 이 temper 을 역연산할 수 있는 untemper 코드가 존재한다. 

Temper

y = mt[self->index++];
y ^= (y >> 11);
y ^= (y << 7) & 0x9d2c5680U;
y ^= (y << 15) & 0xefc60000U;
y ^= (y >> 18);

UnTemper

TemperingMaskB = 0x9d2c5680
TemperingMaskC = 0xefc60000

def untemper(y):
    y = undoTemperShiftL(y)
    y = undoTemperShiftT(y)
    y = undoTemperShiftS(y)
    y = undoTemperShiftU(y)
    return y

def undoTemperShiftL(y):
    last14 = y >> 18
    final = y ^ last14
    return final

def undoTemperShiftT(y):
    first17 = y << 15
    final = y ^ (first17 & TemperingMaskC)
    return final

def undoTemperShiftS(y):
    a = y << 7
    b = y ^ (a & TemperingMaskB)
    c = b << 7
    d = y ^ (c & TemperingMaskB)
    e = d << 7
    f = y ^ (e & TemperingMaskB)
    g = f << 7
    h = y ^ (g & TemperingMaskB)
    i = h << 7
    final = y ^ (i & TemperingMaskB)
    return final

def undoTemperShiftU(y):
    a = y >> 11
    b = y ^ a
    c = b >> 11
    final = y ^ c
    return final

아래는 이에 대한 설명인데, 필자의 지식이 너무나도 부족하여 undoTemperShiftS(y) 함수에 대한 설명은 충분하지 않다.. 나중에 추가할 계획인데, 이에 의문을 가지고 온 사람들에게는 사과의 말씀을....

 

자! 여튼 이렇게 state 를 복구할 수 있겠다. 이를 테스트 해보기 위해서 직접 코드를 만들어서 실행해보자.

import random

output = [random.getrandbits(32) for _ in range(1000)]
init_state = (3,tuple([untemper(op) for op in output[:624]]+[0]),None)

random.setstate(init_state)

for idx in range(1000):
    assert random.getrandbits(32) == output[idx]

- 출처 : https://rbtree.blog/posts/2021-05-18-breaking-python-random-module/

Breaking Python random module

이 경우에, Assertion Error 없이 정상적으로 코드가 진행되는 것을 확인할 수 있다. ( 여기서 init_state 에 들어가는 인자들 중 3 은 version, 그리고 output 의 마지막 0 은 index 를 의미한다. )

 

이로써 32 bit 를 가지고 있을 때, 그 후의 random 값들을 모두 알아낼 수 있다는 것이 증명되었다. 하지만, 문제의 코드를 다시 살펴보면, 32 bit 뿐만 아니라 8, 16 등의 bit 를 알고 있을 때도 이 과정이 통해야 한다. 

 

자!! 그러면 8 bit, 16bit 만 노출되어 있을 때는 어떻게 leak 하는지 알아보는 것은 시험 끝나고 ㅋㅋ

📖 Exploit Code