#!/usr/bin/env python3

import hid
import sys
import os
import zlib
import io
import time
from ctypes import *

from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_PSS
from Crypto.Util.number import bytes_to_long

JEDI_CA_FINGERPRINT = b'\xe5\xe0\x95\xe6C\xb5h\x8b@\x0cu{LD\xef\xac\xc2\x93aH\xe5\xce\xbdlmA\x0fT\xf1H\x7fI'

class DS4Auth(LittleEndianStructure):
    _fields_ = [
        ('type', c_uint8),
        ('seq', c_uint8),
        ('page', c_uint8),
        ('sbz', c_uint8),
        ('data', c_uint8 * 56),
        ('crc32', c_uint32)
    ]

class DS4AuthResult(LittleEndianStructure):
    _fields_ = [
        ('type', c_uint8),
        ('seq', c_uint8),
        ('status', c_uint8),
        ('padding', c_uint8 * 9),
        ('crc32', c_uint32)
    ]

class DS4IdentityBlock(BigEndianStructure):
    _fields_ = [
        ('serial', c_uint8 * 0x10),
        ('modulus', c_uint8 * 0x100),
        ('padding', c_uint8 * 0xfc),
        ('exponent', c_uint32)
    ]

class DS4Response(LittleEndianStructure):
    _fields_ = [
        ('sig_nonce', c_uint8 * 0x100),
        ('identity', DS4IdentityBlock),
        ('sig_identity', c_uint8 * 0x100)
    ]

#DS4_VID = 0x16c0
#DS4_PID = 0x04d5

DS4_VID = 0x054c
DS4_PID = 0x05c4

if __name__ == '__main__':
    dev = hid.device()
    try:
        dev.open(DS4_VID, DS4_PID)
        nonce = io.BytesIO(os.urandom(256))

        print('nonce =', nonce.getvalue().hex())
        challengebuf = DS4Auth()
        challengebuf.seq = 1
        challengebuf.type = 0xf0
        while nonce.readinto(challengebuf.data) != 0:
            print('page =', challengebuf.page, 'data =', memoryview(challengebuf.data).hex())
            challengebuf.crc32 = zlib.crc32(bytes(challengebuf)[:sizeof(DS4Auth)-sizeof(c_uint32)])
            print('crc =', challengebuf.crc32)
            dev.send_feature_report(challengebuf)
            challengebuf.page += 1
            print('sleeping')
            time.sleep(1)
            memset(challengebuf.data, 0, sizeof(challengebuf.data))

        while True:
            begin = time.perf_counter()
            recv = dev.get_feature_report(0xf2, sizeof(DS4AuthResult))
            result = DS4AuthResult()
            size = min(len(recv), sizeof(result))
            memmove(addressof(result), bytes(recv), size)
            print('seq =', result.seq, 'status =', result.status)
            if result.seq != challengebuf.seq:
                print('oops')
                break
            if result.status == 0:
                break
            if time.perf_counter() - begin > 10.0:
                print('timeout waiting for ok')
                break
            time.sleep(1)
        print('auth ok')

        resp = io.BytesIO()

        for i in range(0x13):
            recv = dev.get_feature_report(0xf1, sizeof(DS4Auth))
            respbuf = DS4Auth()
            size = min(len(recv), sizeof(respbuf))
            memmove(addressof(respbuf), bytes(recv), size)
            print('seq =', respbuf.seq, 'page =', respbuf.page)
            print('crc =', respbuf.crc32)
            if zlib.crc32(bytes(respbuf)[:sizeof(DS4Auth)-sizeof(c_uint32)]) != respbuf.crc32:
                print('crc mismatch')
            print('data =', memoryview(respbuf.data).hex())
            resp.write(respbuf.data)

            time.sleep(1)

        print('resp =', resp.getbuffer().hex())

        # verify
        resp.seek(0)
        resp_check = DS4Response()
        resp.readinto(resp_check)
        print('serial =', memoryview(resp_check.identity.serial).hex())

        pss_ca = None
        if os.path.isfile('jedi.pem'):
            with open('jedi.pem', 'rb') as f:
                ca = RSA.importKey(f.read())

            if SHA256.new(ca.exportKey('DER')).digest() != JEDI_CA_FINGERPRINT:
                print('WARNING: Wrong fingerprint for Jedi CA, disabling authenticity check')
            else:
                pss_ca = PKCS1_PSS.new(ca)
        else:
            print('WARNING: Cannot open jedi.pem, disabling authenticity check')

        cuk = RSA.construct((bytes_to_long(bytes(resp_check.identity.modulus)), resp_check.identity.exponent))

        sha_nonce = SHA256.new(nonce.getvalue())
        sha_identity = SHA256.new(resp_check.identity)

        pss_cuk = PKCS1_PSS.new(cuk)


        print('sig_nonce =', memoryview(resp_check.sig_nonce).hex())
        print('identity =', memoryview(resp_check.identity).hex())
        print('sig_identity =', memoryview(resp_check.sig_identity).hex())

        result_sig_nonce = pss_cuk.verify(sha_nonce, bytes(resp_check.sig_nonce))
        if pss_ca is not None:
            result_identity = pss_ca.verify(sha_identity, bytes(resp_check.sig_identity))
        else:
            result_identity = None

        if result_sig_nonce:
            print('good sig for nonce')
        else:
            print('bad sig for nonce')

        if result_identity:
            print('good sig for controller unique key')
        elif result_identity is None:
            print('cannot decide the authenticity of controller unique key')
        else:
            print('bad sig for controller unique key')

        if result_sig_nonce and result_identity:
            print('the controller seems to be genuine')
        elif result_identity is None:
            print('cannot decide the authenticity of the controller')
        else:
            print('the controller may be fake')
    finally:
        dev.close()