Qubic Rube (SECCON 2017 online)

QubicRube.py

#!/bin/python
"""
環境構築メモ
- Python 3.6.3
python -m pip install pyzbar
python -m pip install pillow
python -m pip install opencv-python
python -m pip install scipy

アルゴリズムの説明
1. QRコード(6枚すべて)をスクレイピングでとってくる
2. 各画像を9分割して、背景色ごとに種類分けする
3. 背景色ごとに種類分けした画像をポジションごとに種類分けする(角,端,真ん中)
    - (角,端,真ん中)の特定は、上下左右に黒色がないかで特定できる
4. 角,端,真ん中が判ればQRコードのパターンは 4*4!*4!/4 = 576なので総当たりで殴る!!

計算の合計は 50*6*576 = 172800 で1時間半ぐらいで走り切る
"""

from pyzbar.pyzbar import decode
from PIL import Image
from scipy import ndimage
import urllib.request
import cv2
import itertools
import re

fastHash = '01000000000000000000'
bgr = {
    'yellow' : [0, 213, 255],
    'blue' : [186, 81, 0],
    'green' : [96, 158, 0],
    'white' : [255, 255, 255],
    'red' : [58, 30, 196],
    'orange' : [0, 88, 255],
    'black' : [0, 0, 0],
}

def main():
    hash = fastHash
    #50回回す
    for count in range(1,51):
    # for count in range(1,3):
        print(str(count) + '/50')
        bg_color = {
            'yellow' : [],
            'blue' : [],
            'green' : [],
            'white' : [],
            'red' : [],
            'orange' : [],
        }
        scrap(count, hash)
        hash = None # clear
        for txt in ['U','D','L','R','F','B']:
            fileName = str(count) + '_' + txt + '.png'
        
            input_image = cv2.imread(fileName)
            # cv2.imwrite(tmp_image_name, toBinarization(img))

            # 画像を9等分
            for i in cutNine(input_image):
                c =  getBackColor(i)
                # 背景別に種類分けリストに追加
                if  c == bgr['yellow']:
                    bg_color['yellow'].append(i)
                elif c == bgr['blue']:
                    bg_color['blue'].append(i)
                elif c ==  bgr['green']:
                    bg_color['green'].append(i)
                elif c ==  bgr['white']:
                    bg_color['white'].append(i)
                elif c ==  bgr['red']:
                    bg_color['red'].append(i)
                elif c == bgr['orange']:
                    bg_color['orange'].append(i)

        # 各色ごとにQRを完成させる
        for color, qr_img in bg_color.items():
            break_flag = False

            # 画像をジャンル分け
            qr_parts_corner = []
            qr_parts_side  = []
            qr_parts_center = []

            for m_img in qr_img:
                coordinateType = getCoordinateType(m_img)
                if coordinateType == 'corner':
                    qr_parts_corner.append(m_img)
                elif coordinateType == 'side':
                    qr_parts_side.append(m_img)
                else:
                    qr_parts_center.append(m_img)
            # 4!*4!回画像を組み立てる
            for qr_pattern_corner in list(itertools.permutations( range(4) )):
                for qr_pattern_side in list(itertools.permutations( range(4) )):
                    topImg    = cv2.hconcat( [changeTiltUpperLeft(qr_parts_corner[qr_pattern_corner[0]]),
                                changeTiltTop( qr_parts_side[qr_pattern_side[0]]),
                                changeTiltUpperRight (qr_parts_corner[qr_pattern_corner[2]]) ] )
                    centerImg = cv2.hconcat( [changeTiltLeft(qr_parts_side[qr_pattern_side[1] ]),
                                qr_parts_center[0],
                                changeTiltRight (qr_parts_side[qr_pattern_side[2]]) ])
                    bottomImg = cv2.hconcat( [changeTiltBottomLeft(qr_parts_corner[qr_pattern_corner[1]]),
                                changeTiltBottom( qr_parts_side[qr_pattern_side[3]]),
                                changeTiltBottomRight (qr_parts_corner[qr_pattern_corner[3]]) ] )
                    qr_image = cv2.vconcat([topImg, centerImg, bottomImg])
                    # 背景色を白にして画像を保存
                    cv2.imwrite('qr_image.png', toBinarization(qr_image))
                    # QRコードを読み取り 失敗したら次のパターンを試す
                    try:
                        data = decode(Image.open('qr_image.png'))
                        text = data[0][0].decode('utf-8', 'ignore')
                        print('[-] ' + color + ' : ' + text)
                        break_flag = True
                        if re.match(r"^http" , text):
                            hash = text.split('/')[-1]
                    except IndexError:
                        pass
                    if break_flag:
                        break
                if break_flag:
                    break

# スクレイピング
def scrap(count, hash):
    for txt in ['U','D','L','R','F','B']:
        url = 'http://qubicrube.pwn.seccon.jp:33654/images/' + hash + '_' + txt + '.png'
        fileName = str(count) + '_' + txt + '.png'
        urllib.request.urlretrieve(url, fileName)

# 9等分
def cutNine(img):
    height, width = img.shape[:2]
    m_height = int(height/3)
    m_width  = int(width/3)
    dst1 = img[0 * m_height:1 * m_height, 0 * m_width:1 * m_width]
    dst2 = img[0 * m_height:1 * m_height, 1 * m_width:2 * m_width]
    dst3 = img[0 * m_height:1 * m_height, 2 * m_width:3 * m_width]
    dst4 = img[1 * m_height:2 * m_height, 0 * m_width:1 * m_width]
    dst5 = img[1 * m_height:2 * m_height, 1 * m_width:2 * m_width]
    dst6 = img[1 * m_height:2 * m_height, 2 * m_width:3 * m_width]
    dst7 = img[2 * m_height:3 * m_height, 0 * m_width:1 * m_width]
    dst8 = img[2 * m_height:3 * m_height, 1 * m_width:2 * m_width]
    dst9 = img[2 * m_height:3 * m_height, 2 * m_width:3 * m_width]
    return [dst1, dst2, dst3, dst4, dst5, dst6, dst7, dst8, dst9]

# 背景を全部白色に。
def toBinarization(img):
    height, width = img.shape[:2]
    for i in range(height):
        for j in range(width):
            blue = img.item(i, j, 0)
            green = img.item(i, j, 1)
            red = img.item(i, j, 2)
            if blue != 0 or green != 0 or red != 0:
                img[i, j] = [255, 255, 255]
    return img

# 背景色を返す
def getBackColor(img):
    height, width = img.shape[:2]
    for i in range(height):
        for j in range(width):
            blue = img.item(i, j, 0)
            green = img.item(i, j, 1)
            red = img.item(i, j, 2)
            # if 0 not in [blue, green, red]:
            if blue != 0 or green != 0 or red != 0:
                return [blue, green, red]
    return 0,0,0

"""
+---+---+---+
| 1 | 2 | 3 |
+---+---+---+
| 4 | 5 | 6 |
+---+---+---+
| 7 | 8 | 9 |
+---+---+---+
"""
def getCoordinateId(img):
    height, width = img.shape[:2]
    cutLevel = 6
    imgTop    = img[0:int(height/cutLevel), 0:width]
    imgBottom = img[(cutLevel - 1)*int(height/cutLevel):height, 0:width]
    imgLeft   = img[0:height, 0:int(width/cutLevel)]
    imgRight  = img[0:height, (cutLevel - 1)*int(width/cutLevel):width]

    is_Top = chackAllColor(imgTop)
    is_Bottom = chackAllColor(imgBottom)
    is_Left = chackAllColor(imgLeft)
    is_Right = chackAllColor(imgRight)

    if is_Top and is_Left:
        return 1
    elif is_Bottom and is_Left:
        return 9
    elif is_Top and is_Right:
        return 3
    elif is_Bottom and is_Right:
        return 7
    elif is_Top:
        return 2
    elif is_Bottom:
        return 8
    elif is_Left:
        return 4
    elif is_Right:
        return 6
    else:
        return 5

def getCoordinateType(img):
    coordinateId = getCoordinateId(img)
    if coordinateId in [1, 3, 7, 9]:
        return 'corner'
    elif coordinateId in [2, 4, 6, 8]:
        return 'side'
    else:
        return 'center'

def chackAllColor(img):
    height, width = img.shape[:2]
    for i in range(height):
        for j in range(width):
            blue = img.item(i, j, 0)
            green = img.item(i, j, 1)
            red = img.item(i, j, 2)
            if [blue, green, red] ==  bgr['black']:
                return False
    return True

def tilt90(img):
    img = ndimage.rotate(img, 90, reshape=False)
    return img

# 画像の向きを変える
def changeTiltUpperLeft(img):
    while getCoordinateId(img) != 1:
        img = tilt90(img)
    return img
def changeTiltBottomRight(img):
    while getCoordinateId(img) != 7:
        img = tilt90(img)
    return img
def changeTiltUpperRight(img):
    while getCoordinateId(img) != 3:
        img = tilt90(img)
    return img
def changeTiltBottomLeft(img):
    while getCoordinateId(img) != 9:
        img = tilt90(img)
    return img

def changeTiltTop(img):
    while getCoordinateId(img) != 2:
        img = tilt90(img)
    return img
def changeTiltBottom(img):
    while getCoordinateId(img) != 8:
        img = tilt90(img)
    return img
def changeTiltLeft(img):
    while getCoordinateId(img) != 4:
        img = tilt90(img)
    return img
def changeTiltRight(img):
    while getCoordinateId(img) != 6:
        img = tilt90(img)
    return img

if __name__ == '__main__':
    main()

What is function for??

It kinda like encoder to me