/*
 *      Copyright (C) 2003-2006 Gabest
 *      Copyright (C) 2010-2016 Hendrik Leppkes
 *      http://www.1f0.de
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with this program; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include "stdafx.h"

#include <dvdmedia.h>
#include <ks.h>
#include <ksmedia.h>

#include "DeCSSInputPin.h"
#include "CSSauth.h"
#include "CSSscramble.h"

#include "../DShowUtil.h"

//
// CDeCSSPinHelper
//

CDeCSSPinHelper::CDeCSSPinHelper()
{
  m_varient = -1;
  memset(m_Challenge, 0, sizeof(m_Challenge));
  memset(m_KeyCheck, 0, sizeof(m_KeyCheck));
  memset(m_DiscKey, 0, sizeof(m_DiscKey));
  memset(m_TitleKey, 0, sizeof(m_TitleKey));
}

void CDeCSSPinHelper::Decrypt(IMediaSample* pSample)
{
  long len = pSample->GetActualDataLength();

  BYTE* p = nullptr;
  if(SUCCEEDED(pSample->GetPointer(&p)) && len > 0) {
    if(m_mt.majortype == MEDIATYPE_DVD_ENCRYPTED_PACK && len == 2048 && (p[0x14]&0x30)) {
      CSSdescramble(p, m_TitleKey);
      p[0x14] &= ~0x30;

      IMediaSample2 *pMS2 = nullptr;
      if(SUCCEEDED(pSample->QueryInterface(&pMS2)) && pMS2) {
        AM_SAMPLE2_PROPERTIES props;
        memset(&props, 0, sizeof(props));
        if(SUCCEEDED(pMS2->GetProperties(sizeof(props), (BYTE*)&props))
          && (props.dwTypeSpecificFlags & AM_UseNewCSSKey)) {
            props.dwTypeSpecificFlags &= ~AM_UseNewCSSKey;
            pMS2->SetProperties(sizeof(props), (BYTE*)&props);
        }
        pMS2->Release();
      }
    }
  }
}

void CDeCSSPinHelper::StripPacket(BYTE*& p, long& len)
{
  GUID majortype = m_mt.majortype;
  GUID subtype = m_mt.subtype;

  if(majortype == MEDIATYPE_MPEG2_PACK || majortype == MEDIATYPE_DVD_ENCRYPTED_PACK) {
    if(len > 0 && *(DWORD*)p == 0xba010000) { // MEDIATYPE_*_PACK
      len -= 14;
      p += 14;
      if(int stuffing = (p[-1]&7)) {
        len -= stuffing;
        p += stuffing;
      }
      majortype = MEDIATYPE_MPEG2_PES;
    }
  }

  if(majortype == MEDIATYPE_MPEG2_PES) {
    if(len > 0 && *(DWORD*)p == 0xbb010000) {
      len -= 4;
      p += 4;
      int hdrlen = ((p[0]<<8)|p[1]) + 2;
      len -= hdrlen;
      p += hdrlen;
    }

    if(len > 4
      && ((*(DWORD*)p&0xf0ffffff) == 0xe0010000
      || (*(DWORD*)p&0xe0ffffff) == 0xc0010000
      || (*(DWORD*)p&0xbdffffff) == 0xbd010000)) { // PES
        bool ps1 = (*(DWORD*)p&0xbdffffff) == 0xbd010000;

        len -= 4;
        p += 4;
        size_t expected = ((p[0]<<8)|p[1]);
        len -= 2;
        p += 2;
        BYTE* p0 = p;

        for(int i = 0; i < 16 && *p == 0xff; i++, len--, p++) {
          ;
        }

        if((*p&0xc0) == 0x80) { // mpeg2
          len -= 2;
          p += 2;
          len -= *p+1;
          p += *p+1;
        } else { // mpeg1
          if((*p&0xc0) == 0x40) {
            len -= 2;
            p += 2;
          }

          if((*p&0x30) == 0x30 || (*p&0x30) == 0x20) {
            bool pts = !!(*p&0x20), dts = !!(*p&0x10);
            if(pts) {
              len -= 5;
            }
            p += 5;
            if(dts) {
              ASSERT((*p&0xf0) == 0x10);
              len -= 5;
              p += 5;
            }
          } else {
            len--;
            p++;
          }
        }

        if(ps1) {
          len--;
          p++;
          if(subtype == MEDIASUBTYPE_DVD_LPCM_AUDIO) {
            len -= 6;
            p += 6;
          } else if(subtype == MEDIASUBTYPE_DOLBY_AC3 || subtype == FOURCCMap(0x2000)
            || subtype == MEDIASUBTYPE_DTS || subtype == FOURCCMap(0x2001)) {
              len -= 3;
              p += 3;
          }
        }

        if(expected > 0) {
          expected -= (p - p0);
          len = min((long)expected, len);
        }
    }

    if(len < 0) {
      ASSERT(0);
      len = 0;
    }
  }
}

// IKsPropertySet

STDMETHODIMP CDeCSSPinHelper::Set(REFGUID PropSet, ULONG Id, LPVOID pInstanceData, ULONG InstanceLength, LPVOID pPropertyData, ULONG DataLength)
{
  if(PropSet != AM_KSPROPSETID_CopyProt) {
    return E_NOTIMPL;
  }

  switch(Id) {
  case AM_PROPERTY_COPY_MACROVISION:
    break;
  case AM_PROPERTY_DVDCOPY_CHLG_KEY: 
    { // 3. auth: receive drive nonce word, also store and encrypt the buskey made up of the two nonce words
      AM_DVDCOPY_CHLGKEY* pChlgKey = (AM_DVDCOPY_CHLGKEY*)pPropertyData;
      for(int i = 0; i < 10; i++) {
        m_Challenge[i] = pChlgKey->ChlgKey[9-i];
      }

      CSSkey2(m_varient, m_Challenge, &m_Key[5]);

      CSSbuskey(m_varient, m_Key, m_KeyCheck);
    }
    break;
  case AM_PROPERTY_DVDCOPY_DISC_KEY: 
    { // 5. receive the disckey
      AM_DVDCOPY_DISCKEY* pDiscKey = (AM_DVDCOPY_DISCKEY*)pPropertyData; // pDiscKey->DiscKey holds the disckey encrypted with itself and the 408 disckeys encrypted with the playerkeys

      bool fSuccess = false;

      for(int j = 0; j < g_nPlayerKeys; j++) {
        for(int k = 1; k < 409; k++) {
          BYTE DiscKey[6];
          for(int i = 0; i < 5; i++) {
            DiscKey[i] = pDiscKey->DiscKey[k*5+i] ^ m_KeyCheck[4-i];
          }
          DiscKey[5] = 0;

          CSSdisckey(DiscKey, g_PlayerKeys[j]);

          BYTE Hash[6];
          for(int i = 0; i < 5; i++) {
            Hash[i] = pDiscKey->DiscKey[i] ^ m_KeyCheck[4-i];
          }
          Hash[5] = 0;

          CSSdisckey(Hash, DiscKey);

          if(!memcmp(Hash, DiscKey, 6)) {
            memcpy(m_DiscKey, DiscKey, 6);
            j = g_nPlayerKeys;
            fSuccess = true;
            break;
          }
        }
      }

      if(!fSuccess) {
        return E_FAIL;
      }
    }
    break;
  case AM_PROPERTY_DVDCOPY_DVD_KEY1:
    { // 2. auth: receive our drive-encrypted nonce word and decrypt it for verification
      AM_DVDCOPY_BUSKEY* pKey1 = (AM_DVDCOPY_BUSKEY*)pPropertyData;
      for(int i = 0; i < 5; i++) {
        m_Key[i] =  pKey1->BusKey[4-i];
      }

      m_varient = -1;

      for(int i = 31; i >= 0; i--) {
        CSSkey1(i, m_Challenge, m_KeyCheck);

        if(memcmp(m_KeyCheck, &m_Key[0], 5) == 0) {
          m_varient = i;
        }
      }
    }
    break;
  case AM_PROPERTY_DVDCOPY_REGION:
    break;
  case AM_PROPERTY_DVDCOPY_SET_COPY_STATE:
    break;
  case AM_PROPERTY_DVDCOPY_TITLE_KEY:
    { // 6. receive the title key and decrypt it with the disc key
      AM_DVDCOPY_TITLEKEY* pTitleKey = (AM_DVDCOPY_TITLEKEY*)pPropertyData;
      for(int i = 0; i < 5; i++) {
        m_TitleKey[i] = pTitleKey->TitleKey[i] ^ m_KeyCheck[4-i];
      }
      m_TitleKey[5] = 0;
      CSStitlekey(m_TitleKey, m_DiscKey);
    }
    break;
  default:
    return E_PROP_ID_UNSUPPORTED;
  }

  return S_OK;
}

STDMETHODIMP CDeCSSPinHelper::Get(REFGUID PropSet, ULONG Id, LPVOID pInstanceData, ULONG InstanceLength, LPVOID pPropertyData, ULONG DataLength, ULONG* pBytesReturned)
{
  if(PropSet != AM_KSPROPSETID_CopyProt) {
    return E_NOTIMPL;
  }

  switch(Id) {
  case AM_PROPERTY_DVDCOPY_CHLG_KEY: 
    { // 1. auth: send our nonce word
      AM_DVDCOPY_CHLGKEY* pChlgKey = (AM_DVDCOPY_CHLGKEY*)pPropertyData;
      for(int i = 0; i < 10; i++) {
        pChlgKey->ChlgKey[i] = 9 - (m_Challenge[i] = i);
      }
      *pBytesReturned = sizeof(AM_DVDCOPY_CHLGKEY);
    }
    break;
  case AM_PROPERTY_DVDCOPY_DEC_KEY2:
    { // 4. auth: send back the encrypted drive nonce word to finish the authentication
      AM_DVDCOPY_BUSKEY* pKey2 = (AM_DVDCOPY_BUSKEY*)pPropertyData;
      for(int i = 0; i < 5; i++) {
        pKey2->BusKey[4-i] = m_Key[5+i];
      }
      *pBytesReturned = sizeof(AM_DVDCOPY_BUSKEY);
    }
    break;
  case AM_PROPERTY_DVDCOPY_REGION: 
    {
      DVD_REGION* pRegion = (DVD_REGION*)pPropertyData;
      pRegion->RegionData = 0;
      pRegion->SystemRegion = 0;
      *pBytesReturned = sizeof(DVD_REGION);
    }
    break;
  case AM_PROPERTY_DVDCOPY_SET_COPY_STATE: 
    {
      AM_DVDCOPY_SET_COPY_STATE* pState = (AM_DVDCOPY_SET_COPY_STATE*)pPropertyData;
      pState->DVDCopyState = AM_DVDCOPYSTATE_AUTHENTICATION_REQUIRED;
      *pBytesReturned = sizeof(AM_DVDCOPY_SET_COPY_STATE);
    }
    break;
  default:
    return E_PROP_ID_UNSUPPORTED;
  }

  return S_OK;
}

STDMETHODIMP CDeCSSPinHelper::QuerySupported(REFGUID PropSet, ULONG Id, ULONG* pTypeSupport)
{
  if(PropSet != AM_KSPROPSETID_CopyProt) {
    return E_NOTIMPL;
  }

  switch(Id) {
  case AM_PROPERTY_COPY_MACROVISION:
    *pTypeSupport = KSPROPERTY_SUPPORT_SET;
    break;
  case AM_PROPERTY_DVDCOPY_CHLG_KEY:
    *pTypeSupport = KSPROPERTY_SUPPORT_GET | KSPROPERTY_SUPPORT_SET;
    break;
  case AM_PROPERTY_DVDCOPY_DEC_KEY2:
    *pTypeSupport = KSPROPERTY_SUPPORT_GET;
    break;
  case AM_PROPERTY_DVDCOPY_DISC_KEY:
    *pTypeSupport = KSPROPERTY_SUPPORT_SET;
    break;
  case AM_PROPERTY_DVDCOPY_DVD_KEY1:
    *pTypeSupport = KSPROPERTY_SUPPORT_SET;
    break;
  case AM_PROPERTY_DVDCOPY_REGION:
    *pTypeSupport = KSPROPERTY_SUPPORT_GET | KSPROPERTY_SUPPORT_SET;
    break;
  case AM_PROPERTY_DVDCOPY_SET_COPY_STATE:
    *pTypeSupport = KSPROPERTY_SUPPORT_GET | KSPROPERTY_SUPPORT_SET;
    break;
  case AM_PROPERTY_DVDCOPY_TITLE_KEY:
    *pTypeSupport = KSPROPERTY_SUPPORT_SET;
    break;
  default:
    return E_PROP_ID_UNSUPPORTED;
  }

  return S_OK;
}

//
// CDeCSSTransformInputPin
//

CDeCSSTransformInputPin::CDeCSSTransformInputPin(TCHAR* pObjectName, CTransformFilter* pFilter, HRESULT* phr, LPWSTR pName)
  : CTransformInputPin(pObjectName, pFilter, phr, pName)
  , CDeCSSPinHelper()
{
}

STDMETHODIMP CDeCSSTransformInputPin::NonDelegatingQueryInterface(REFIID riid, void** ppv)
{
  CheckPointer(ppv, E_POINTER);

  return
    QI(IKsPropertySet)
    __super::NonDelegatingQueryInterface(riid, ppv);
}

STDMETHODIMP CDeCSSTransformInputPin::Receive(IMediaSample* pSample)
{
  Decrypt(pSample);
  return __super::Receive(pSample);
}

HRESULT CDeCSSTransformInputPin::SetMediaType(const CMediaType *pmt)
{
  SetCSSMediaType(pmt);
  return __super::SetMediaType(pmt);
}