#include <windows.h>
#include <winioctl.h>
#include <stdio.h>
#include "ntddcdvd.h"

#include "CSSauth.h"

HANDLE hDVD;
DVD_SESSION_ID Session;

BYTE Challenge[10];
BYTE Key1[5];
BYTE Key2[5];
BYTE KeyBuffer[DVD_DISK_KEY_LENGTH];
BYTE KeyCheck[5];

typedef struct
{
	BYTE lchal[10];
	BYTE hchal[10];
	BYTE lkey1[5];
	BYTE hkey2[5];
} DVD_AUTHINFO;

HANDLE DVDOpen(char drive)
{
	HANDLE h;
	char temp[100];
	sprintf(temp,"\\\\.\\%c:",drive);
	h=CreateFile(temp,GENERIC_READ|GENERIC_WRITE,FILE_SHARE_READ|FILE_SHARE_WRITE,NULL,OPEN_EXISTING,0,NULL);
	return(h);
}

void DVDClose(HANDLE h)
{
	CloseHandle(h);
}

DVD_SESSION_ID DVDStartSession(HANDLE h)
{
	DWORD r;
	ULONG id;
	if(!DeviceIoControl(h,IOCTL_DVD_START_SESSION,&r,4,&id,sizeof(id),&r,NULL))
		return(0);
	return(id);
}

void DVDEndSession(HANDLE h,DVD_SESSION_ID id)
{
	DWORD r;
	DeviceIoControl(h,IOCTL_DVD_END_SESSION,&id,sizeof(id),NULL,0,&r,NULL);
}	


int DVDSendKey(HANDLE h,ULONG s,DVD_KEY_TYPE kt,BYTE *kb)
{
	DWORD r=0,bl,kl;
	PDVD_COPY_PROTECT_KEY key=(PDVD_COPY_PROTECT_KEY)KeyBuffer;
	switch(kt)
	{
		case DvdChallengeKey:
			kl=10;
			bl=12;
			break;
		case DvdBusKey1:
		case DvdBusKey2:
			kl=5;
			bl=8;
			break;
		case DvdDiskKey:
			kl=2048;
			bl=2048;
			break;
		case DvdTitleKey:
			kl=5;
			bl=8;
			break;
		default:
			return(1);
			break;
	}
	key->KeyLength=bl+sizeof(DVD_COPY_PROTECT_KEY);
	key->SessionId=s;
	key->KeyType=kt;
	key->KeyFlags=0;
	memcpy(key->KeyData,kb,kl);
	if(!DeviceIoControl(h,IOCTL_DVD_SEND_KEY,key,key->KeyLength,key,key->KeyLength,&r,NULL))
		return(2);
	return(0);
}

int DVDReadKey(HANDLE h,ULONG s,DVD_KEY_TYPE kt,BYTE *kb,void *data)
{
	DWORD r=0,bl,kl;
	PDVD_COPY_PROTECT_KEY key=(PDVD_COPY_PROTECT_KEY)KeyBuffer;
	switch(kt)
	{
		case DvdChallengeKey:
			kl=10;
			bl=12;
			break;
		case DvdBusKey1:
		case DvdBusKey2:
			kl=5;
			bl=8;
			break;
		case DvdDiskKey:
			kl=2048;
			bl=2048;
			break;
		case DvdTitleKey:
			kl=6;
			bl=8;
			key->Parameters.TitleOffset.QuadPart=*((LONGLONG *)data);
			break;
		default:
			return(1);
			break;
	}
	key->KeyLength=bl+sizeof(DVD_COPY_PROTECT_KEY);
	key->SessionId=s;
	key->KeyType=kt;
	key->KeyFlags=0;
	if(!DeviceIoControl(h,IOCTL_DVD_READ_KEY,key,key->KeyLength,key,key->KeyLength,&r,NULL))
		return(2);
	memcpy(kb,key->KeyData,kl);
	return(0);
}

int CSSauthenticate(HANDLE hDVD)
{
	int title=0;
	int index,varient;
	DVD_AUTHINFO ai;

	for(index=0;index<10;index++)
		Challenge[index]=index;

	for(index=0;index<10;index++)
		ai.hchal[index]=Challenge[9-index];
	if(DVDSendKey(hDVD,Session,DvdChallengeKey,ai.hchal))
		return(1);

	if(DVDReadKey(hDVD,Session,DvdBusKey1,ai.lkey1,NULL))
		return(2);
	for(index=0;index<5;index++)
		Key1[index]=ai.lkey1[4-index];

	for(varient=-1,index=31;index>=0;index--)
	{
		CSSkey1(index,Challenge,KeyCheck);
		if(memcmp(KeyCheck,Key1,5)==0)
		{
			varient=index;
		}
	}
	if(DVDReadKey(hDVD,Session,DvdChallengeKey,ai.lchal,NULL))
		return(3);
	for(index=0;index<10;index++)
		Challenge[index]=ai.lchal[9-index];

	CSSkey2(varient,Challenge,Key2);
	for(index=0;index<5;index++)
		ai.hkey2[4-index]=Key2[index];
	if(DVDSendKey(hDVD,Session,DvdBusKey2,ai.hkey2))
		return(1);

	memcpy(Challenge,Key1,5);
	memcpy(Challenge+5,Key2,5);
	CSSbuskey(varient,Challenge,KeyCheck);
	return(0);
}

int CSSgetdiskkey(char n,unsigned char *key)
{
	int i,e;
	BYTE tmp[2048];

	if((hDVD=DVDOpen(n))==INVALID_HANDLE_VALUE)
		return(1);

	if(!(Session=DVDStartSession(hDVD)))
	{
		DVDClose(hDVD);
		return(2);
	}

	e=CSSauthenticate(hDVD);
	if(!e)
	{
		e=DVDReadKey(hDVD,Session,DvdDiskKey,tmp,NULL);
		if(!e)
			for(i=0x235;i<0x23a;i++)
				key[i-0x235]=tmp[i]^KeyCheck[4-(i%5)];
		key[5]=0;
	}
	if(Session)
		DVDEndSession(hDVD,Session);
	DVDClose(hDVD);
	return(e);
}

int CSSgettitlekey(char n,unsigned int LBA,unsigned char *key)
{
	int i,e;
	LONGLONG t=LBA*2048;

	if((hDVD=DVDOpen(n))==INVALID_HANDLE_VALUE)
		return(1);

	if(!(Session=DVDStartSession(hDVD)))
	{
		DVDClose(hDVD);
		return(2);
	}

	e=CSSauthenticate(hDVD);
	if(!e)
	{
		e=DVDReadKey(hDVD,Session,DvdTitleKey,key,&t);
		if(!e)
			for (i=0;i<5;i++)
				key[i]^=KeyCheck[4-(i%5)];
		key[5]=0;
	}
	if(Session)
		DVDEndSession(hDVD,Session);
	DVDClose(hDVD);
	return(e);
}