#!perl -w
# (C) 2003-2007 Willem Jan Hengeveld <itsme@xs4all.nl>
# Web: http://www.xs4all.nl/~itsme/
#      http://wiki.xda-developers.com/
#
# $Id$
#
use strict;
use IO::File;
use Win32API::File qw( :ALL );
use Win32API::Registry qw( :ALL );
use Dumpvalue;
my $d= new Dumpvalue;

use Getopt::Long;

sub usage {
return <<__EOF__;
Usage: xdarit [-l] [-i] [-x] [-w] [-k] [--ce=<cerom> ] [--bl=<bootloader> ] [--gsm=<gsmrom> ]
   -l  : list disks
   -i  : print info about sd/mmc-disk
   -r  : read image(s) from sd/mmc-disk
   -w  : write image(s) to sd/mmc-disk
   -k  : add flashkey when writing ( for bootloader >= 5.17 )
   -a  : extract all images
   -d N: specify disk - defaults to first found removable
__EOF__
}

my @actions;
my %params;
GetOptions(
    "listdisks|l" => sub { push @actions, \&ActionListDisks; },
    "info|i" => sub { push @actions, \&ActionDumpSDInfo; },
    "ce=s" => sub { $params{ceimagename}= $_[1]; },
    "bl=s" => sub { $params{blimagename}= $_[1]; },
    "gsm=s" => sub { $params{gsmimagename}= $_[1]; },
    "all|a" => sub { $params{allimages}= 1; },
    "write|w" => sub { push @actions, \&ActionWriteImagesToSD; },
    "read|r" => sub { push @actions, \&ActionReadImagesFromSD; },
    "key|k" => sub { $params{addkey}= 1; },
    "disk|d=s" => sub { $params{disknr}= eval($_[1]); },
) or die usage();

die usage() if (!@actions);

#GetStorageId("\\\\.\\$_:") for ('I');
#GetStorageId("\\\\.\\PhysicalDrive$_") for (1);
for (@actions) {
    $_->(%params);
}
exit(0);
###############################################################################
##  actions
sub ActionListDisks {
    # print registry key info
    for my $i (0..GetNrOfDisks()-1)
    {
        my $info=GetDiskRegInfo($i);
        my ($drvid)= ($info->{Driver} =~ /\\(\w+)$/);
        printf("%2d %s  %s\n", $i, $drvid, $info->{FriendlyName});
# join("", map { "      $_\n" } @{$info->{HardwareID}}), $info->{EnumKey}
    }

    # 
    for my $disk (GetMediaList()) {
        my $data= substr(unpack("A*", $disk->{first}),0, 32);
        $data =~ s/\W/ /g;
        my $dosname;
        QueryDosDevice("PhysicalDrive$disk->{index}", $dosname, '1024') or warn "$^E\n";
        printf("%2d %8s '%s' %s\n", $disk->{index}, 
            CapacityString($disk->{BytesPerSector} * $disk->{SectorsPerTrack} * $disk->{TracksPerCylinder} * $disk->{Cylinders}),
            $data, $dosname);
    }

    my $drvstrings;
    GetLogicalDriveStrings('1000', $drvstrings) or warn "$^E\n";
    for (split /\0/, $drvstrings)
    {
        my ($volname, $serialnum, $maxpathlen, $flags, $fstype);
        if (!GetVolumeInformation($_, $volname, '1024', $serialnum, $maxpathlen, $flags, $fstype, '1024'))
        {
            printf("%s : %s\n", $_, $^E);
        }
        else {
            printf("%s : %s %08lx %08lx %08lx %s\n",
                $_, $volname, $serialnum, $maxpathlen, $flags, $fstype);
        }
    }
}
sub CapacityString {
    my ($bytes)= @_;

    if ($bytes<100) { return sprintf("%d", $bytes); }

    my @letters= qw(0 K M G T P E);
    my @formats= qw(%5.2f %5.1f %5.1f);

    for (my $i= 1 ; $i<=@formats*@letters ; $i++ ) {
        my $limit= 10**$i;
        my $size= 10**($i-($i%3));
        if ($bytes<$limit) { return sprintf("$formats[$i%3]$letters[$i/3]", $bytes/$size); }
    }

    return sprintf("%5.1g", $bytes);
}

sub IDE_ATAPI_IDENTIFY  { 0xA1; }
sub IDE_ATA_IDENTIFY    { 0xEC; }
sub SMART_GET_VERSION          { 0x00074080; }
sub SMART_SEND_DRIVE_COMMAND   { 0x0007c084; }
sub SMART_RECEIVE_DRIVE_DATA   { 0x0007c088; }
sub IDENTIFY_BUFFER_SIZE  { 512; }

sub ActionDumpSDInfo {
    my (%params)= @_;

    my $disk= GetSpecifiedDisk(%params);
    print "-------------disk-------------------\n";
    $d->dumpValue($disk);

    print "-------------diskinfo---------------\n";
    my $diskinfo= GetSDInfo($disk) if ($disk->{BytesPerSector});
    $d->dumpValue($diskinfo);

    GetStorageId($disk->{devicename});

    print "-------------driveinfo---------------\n";
    my %driveinfo;
    if (GetDriveVersion($disk->{devicename}, \%driveinfo)) {
        $d->dumpValue(\%driveinfo);
    }
    my $head= 0xa0 |(($disk->{index}&1)<<4);
    my $cmd= ($driveinfo{IDEDeviceMap}>>$disk->{index})&0x10 ? IDE_ATAPI_IDENTIFY : IDE_ATA_IDENTIFY;

    print "-------------drivedata---------------\n";
    my %drivedata;
    GetDriveData($disk->{devicename}, \%drivedata, $head, $cmd, $disk->{index});
    $d->dumpValue(\%drivedata);

}
sub GetSpecifiedDisk {
    my (%params)= @_;
    my @media= GetMediaList() or die "No disks found\n";

    my %disk2media= map { $_->{index} => $_ } @media;

    if (!@media) {
        warn "no media found\n";
        return;
    }
    my $disknr= exists $params{disknr}?$params{disknr}:$media[0]{index};

    return $disk2media{$disknr};
}
sub GetSDInfo {
    my ($disk)= @_;
    my $header= ReadRaw($disk->{devicename}, 0, 7*$disk->{BytesPerSector});

    my %info;

    if (substr($header, 0, 16) eq pack("a16", "HTC FLASH KEY"))
    {
        $info{haskey}= 1;
        $info{headersize}= 7*$disk->{BytesPerSector};
        $header= substr($header, 5*$disk->{BytesPerSector});

        $info{cidmd5}= ReadRaw($disk->{devicename}, $disk->{BytesPerSector}, 16);
        $info{counter}= unpack("L", ReadRaw($disk->{devicename}, 2*$disk->{BytesPerSector}, 4));
        $info{countermd5}= ReadRaw($disk->{devicename}, 3*$disk->{BytesPerSector}, 16);
    }
    else {
        $info{headersize}= 2*$disk->{BytesPerSector};
    }

    # todo: this is cfg info, move to separate function
    if (substr($header, 0, 16) eq pack("a16", 'HTC$WALLABY00')) {
        $info{bootloadersize}= 0x40000;
    }
    elsif (substr($header, 0, 16) eq pack("a16", 'HTC$WALLABY11')) {
        $info{cesize}= 0x01fc0000;
    }
    elsif (substr($header, 0, 16) eq pack("a16", 'HTC$WALLABY22')) {
        $info{bootloadersize}= 0x40000;
        $info{cesize}= 0x01fc0000;
    }
    elsif (substr($header, 0, 16) eq pack("a16", 'HTC$WALLABY33')) {
        $info{diagnosticssize}= 6144;   # ??? just guessing.
        $info{headersize} -= $disk->{BytesPerSector};
    }
    elsif (substr($header, 0, 16) eq pack("a16", 'HTC$WALLABY44')) {
        $info{hasgsmrom}= 1;
    }
    elsif (substr($header, 0, 16) eq pack("a16", 'HTC$WALLABY55')) {
        $info{hasgsmrom}= 1;
        $info{cesize}= 0x01fc0000;
    }

    return \%info;
}
sub ActionWriteImagesToSD {
    my (%params)= @_;

    my $disk= GetSpecifiedDisk(%params);

    my $diskinfo= GetSDInfo($disk);

    my @cfg;# diag, gsm, ce, boot
    $cfg[0][0][0][1]= 'HTC$WALLABY00';
    $cfg[0][0][1][0]= 'HTC$WALLABY11';
    $cfg[0][0][1][1]= 'HTC$WALLABY22';
    $cfg[1][0][0][0]= 'HTC$WALLABY33';
    $cfg[0][1][0][0]= 'HTC$WALLABY44';
    $cfg[0][0][1][0]= 'HTC$WALLABY55';

    # todo: this is wrong, check image names iso sizes.
    my $header= $cfg[$diskinfo->{diagsize}!=0][$diskinfo->{gsmsize}!=0][$diskinfo->{cesize}!=0][$diskinfo->{bootloadersize}!=0];
    if (!$header) {
        die "invalid combinatian of images\n";
    }
    my $ofs= 0;

    # write simple header.
    RawWrite($disk->{devicename}, $ofs, pack("a16", $header));
    $ofs += $disk->{headersize};

    if ($diskinfo->{bootloadersize}) {
        my $bootloader= ReadImage($params{blimagename});
        RawWrite($disk->{devicename}, $ofs, $bootloader);
        $ofs += $disk->{bootloadersize};
    }
    if ($diskinfo->{cesize}) {
        my $ce= ReadImage($params{ceimagename});
        RawWrite($disk->{devicename}, $ofs, $ce);
        $ofs += $disk->{cesize};
    }
    if ($diskinfo->{diagsize}) {
        my $diag= ReadImage($params{diagimagename});
        RawWrite($disk->{devicename}, $ofs, $diag);
        $ofs += $disk->{diagsize};
    }

}
sub ActionReadImagesFromSD {
    # read 7 headerblocks, followed by bootloader, followed by ce-rom
}
###############################################################################
##             low level disk access
#
sub GetRemovableMediaList {
    my @media;
    for my $i (0..GetNrOfDisks()-1) {
        my %info;
        if (ReadDiskInfo($i, \%info)) {
            push @media, \%info if ($info{MediaType} && $info{MediaType}==RemovableMedia);
        }
    }
    return @media;
}
sub GetMediaList {
    my @media;
    for my $i (0..16) {
        my %info;
        if (ReadDiskInfo($i, \%info)) {
            push @media, \%info;
        }
    }
    return @media;
}


# see HKLM\SYSTEM\CurrentControlSet\Services\Disk\Enum
sub GetNrOfDisks {
    my $hkey;
    RegOpenKeyEx(HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\Disk\\Enum", 0, KEY_READ, $hkey )
        or die "Can't open Services\\disk\\Enum key ", regLastError(),"\n";

    my $data; my $type;
    RegQueryValueEx($hkey, "Count", [], $type, $data, [] )
        or die "Cannot find Services\\Disk\\Count key ", regLastError(),"\n";

    RegCloseKey($hkey);
    return unpack("L", $data);
}
sub GetDiskName {
    my ($disknr)= @_;
    my $hkey;
    RegOpenKeyEx(HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\Disk\\Enum", 0, KEY_READ, $hkey )
        or die "Can't open Services\\disk\\Enum key ", regLastError(),"\n";

    my $data; my $type;
    RegQueryValueEx($hkey, "$disknr", [], $type, $data, [] )
        or return !warn "Cannot find Services\\Disk\\Enum\\$disknr key ", regLastError(),"\n";

    RegCloseKey($hkey);

    RegOpenKeyEx(HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Enum\\$data", 0, KEY_READ, $hkey )
        or  die "Can't open Enum\\$data key ", regLastError(),"\n";
    my $name;
    RegQueryValueEx($hkey, "FriendlyName", [], $type, $name, [] )
        or die "Cannot find Enum\\$data\\FriendlyName key ", regLastError(),"\n";

    RegCloseKey($hkey);

    return $name;
}
sub GetDiskRegInfo {
    my ($disknr)= @_;
    my $hkey;
    RegOpenKeyEx(HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Services\\Disk\\Enum", 0, KEY_READ, $hkey )
        or die "Can't open Services\\disk\\Enum key ", regLastError(),"\n";

    my $enumkey; my $type;
    RegQueryValueEx($hkey, "$disknr", [], $type, $enumkey, [] )
        or die "Cannot find Services\\Disk\\Enum\\$disknr key ", regLastError(),"\n";

    RegCloseKey($hkey);

    RegOpenKeyEx(HKEY_LOCAL_MACHINE, "SYSTEM\\CurrentControlSet\\Enum\\$enumkey", 0, KEY_READ, $hkey )
        or  die "Can't open Enum\\$enumkey key ", regLastError(),"\n";

    my %info;
    for (qw(FriendlyName HardwareID Driver)) {
        my $data="";
        RegQueryValueEx($hkey, $_, [], $type, $data, [] )
            or die "Cannot find Enum\\$enumkey\\$_ key ", regLastError(),"\n";
        if ($type==REG_SZ) {
            $info{$_}= $data;
        }
        elsif ($type==REG_MULTI_SZ) {
            $info{$_}= [ split /\0/, $data ];
        }
        elsif ($type==REG_DWORD) {
            $info{$_}= unpack("L", $data);
        }
        else {
            $info{$_}= { type=>$type, data=>$data };
        }
    }

    RegCloseKey($hkey);

    $info{EnumKey}= $enumkey;

    return \%info;
}
sub GetStorageId {
    my ($device)= @_;
# http://msdn.microsoft.com/library/default.asp?url=/library/en-us/dnwmt/html/wmdriverreq.asp
    my $fh= createFile($device, "r", "rw") or return !warn "creatFile($device): $^E\n";

    my $diskid;
    DeviceIoControl($fh, 0x002d0c10, [], 0, $diskid, 4*4+32, [], [])  or return !warn "DeviceIoControl($device, GETID) : $^E\n";

    CloseHandle($fh);

    my @fields= unpack("L4A*", $diskid);
    printf("len=%08lx res=%08lx auth=%08lx res=%08lx  nr=%s\n", @fields);
}

sub GetDriveGeometry {
    my ($device, $info)= @_;

    my $fh= createFile($device, "r", "rw") or return !warn "creatFile($device): $^E\n";

    my $diskgeometry;
    DeviceIoControl($fh, IOCTL_DISK_GET_DRIVE_GEOMETRY, [], 0, $diskgeometry, 6*4, [], [])  or return !warn "DeviceIoControl($device, GETGEOM) : $^E\n";

    CloseHandle($fh);

    my @fields= unpack("L6", $diskgeometry);
    AddToHash($info, {
        Cylinders=>$fields[0]+$fields[1]*(2**32),
        MediaType=>$fields[2],
        TracksPerCylinder=>$fields[3],
        SectorsPerTrack=>$fields[4],
        BytesPerSector=>$fields[5],
    });
    return 1;
}

sub GetDriveVersion {
    my ($device, $info)= @_;
    my $fh= createFile($device, "r", "rw") or return !warn "creatFile($device): $^E\n";
    my $version;
    DeviceIoControl($fh, SMART_GET_VERSION, [], 0, $version, 6*4, [], []) 
        or return !warn "DeviceIoControl($device, GETVERSION) : $^E\n";

    CloseHandle($fh);
    my @fields= unpack("CCCCV5", $version);
    AddToHash($info, {
        Version=>$fields[0],
        Revision=>$fields[1],
        Reserved1=>$fields[2],
        IDEDeviceMap=>$fields[3],
        Capabilities=>$fields[4],
        Reserved2=>[@fields[5..8]],
    });
    return 1;
}
sub GetDriveData {
    my ($device, $info, $head, $command, $drivenum)= @_;
    my $fh= createFile($device, "rw ke", "rw") or return !warn "dd creatFile($device): $^E\n";

    my $data;
    my $scip= pack("VC8CC3V4",
        IDENTIFY_BUFFER_SIZE,
        0,1,1,0,0, $head, $command, 0,
        $drivenum,
        0,0,0,
        0,0,0,0);
    DeviceIoControl($fh, SMART_RECEIVE_DRIVE_DATA, $scip, 0, $data, 4*4+IDENTIFY_BUFFER_SIZE, [], []) 
        or return !warn "DeviceIoControl($device, DRIVEDATA) : $^E\n";

    CloseHandle($fh);
    my @fields= unpack("VCCC2V2a*", $data);
    AddToHash($info, {
        BufferSize=>$fields[0],
        DriverError=>$fields[1],
        IDEError=>$fields[2],
        Reserved1=>[@fields[3..4]],
        Reserved2=>[@fields[5..6]],
        data=>$fields[7],
    });
    return 1;
}
sub ReadRaw {
    my ($device, $offset, $size)= @_;

    my $fh= createFile($device, "r", "rw") or return !warn "creatFile($device): $^E\n";

    SetFilePointer($fh, $offset, [], FILE_BEGIN) or return !warn "SetFilePointer($device): $^E\n";

    my $buffer;
    ReadFile($fh, $buffer, $size, [], []) or return !warn "readfile($device, $size): $^E\n";

    CloseHandle($fh);

    return $buffer;
}
sub ReadDiskInfo {
    my ($index, $info)= @_;

    my $devicename= "\\\\.\\PhysicalDrive$index";
    my $friendlyname= GetDiskName($index);
    my %geometry;
    GetDriveGeometry($devicename, \%geometry) or return;

    my $sector= ReadRaw($devicename, 0, $geometry{BytesPerSector});
    AddToHash($info, {
        index=>$index,
        devicename=>$devicename,
        friendlyname=>$friendlyname,
        %geometry,
        first=>$sector,
    });
    return 1;
}

sub AddToHash {
    my ($hash, $add)= @_;
    $hash->{$_}= $add->{$_} for keys %$add;
}
=pod

card with keysignature have the following structure:

block0: starts with string "HTC FLASH KEY\0"
block1: starts with 16 byte md5 sum of hex representation of first 16 bytes of cardid
block2: start with 4 byte value
block3: starts with md5sum of hex representation of first 4 bytes of block2

=cut