unit Main;

interface

uses Classes, SysUtils, Struct, MyTypes, MainInfo;

type
  EMessage = class(Exception); // Exception raised by our code, The Message field (string)
                               // holds all necessary information, No need to display
                               // the exception class, ErrorMessage field, Win32 last error,
                               // or (try to) track down error address.

  EFatalError = class(Exception); // Halt program. Only handled by the top-level handler,
                                  // not by local handlers. Typically raised when an error
                                  // has been caught and processed but the execution can not
                                  // continue.

procedure InternalError(const Id: String);
{$IFDEF UNICODE}
function TestPassword(Info: TMainInfo; const Password: string): Boolean; overload;
{$ENDIF}
function TestPassword(Info: TMainInfo; const Password: AnsiString): Boolean; {$IFDEF UNICODE} overload; {$ENDIF}
function TestPasswordIs64(Info: TMainInfo; const EncryptionKey: TSetupEncryptionKey): Boolean;
function GetSliceName(Info: TMainInfo; ASlice:integer):string;
function FixPasswordEncoding(Info: TMainInfo; const Password: string) : AnsiString;
procedure GenerateEncryptionKey(const Password: String; const Salt: TSetupKDFSalt; const Iterations: Integer; out Key: TSetupEncryptionKey);

function StrToOem(const s:string):AnsiString;
function OemToStr(const s:AnsiString):string;
procedure write(Info: TMainInfo; s:string);
procedure writeln(Info: TMainInfo; const s:string);

function AddFakeFile(Info: TMainInfo; const FileName: string; const FileContents:AnsiString; RenameNow:boolean=false):integer;

function FileContents(Filename:string) : AnsiString;
function MakeDir(Dir: String) : Boolean;

function GetCustomMessage(Info: TMainInfo; const AText: string): string;

implementation

uses Windows, MD5, SHA1, SHA256, ChaCha20, PBKDF2, PathFunc, CmnFunc2, Msgs, MsgIds, StrUtils;

procedure InternalError(const Id: String);
begin
  raise Exception.Create(Id);
end;

function TestPassword_MD5(const Password: AnsiString; const PasswordHash: TSetupHash; const PasswordSalt: TSetupSalt): Boolean;
var
  Context: TMD5Context;
  Hash: TMD5Digest;
begin
  MD5Init(Context);
  MD5Update(Context, PAnsiChar('PasswordCheckHash')^, Length('PasswordCheckHash'));
  MD5Update(Context, PasswordSalt, SizeOf(PasswordSalt));
  MD5Update(Context, Pointer(Password)^, Length(Password));
  Hash := MD5Final(Context);
  Result := MD5DigestsEqual(Hash, PasswordHash.MD5);
end;

function TestPassword_SHA1(const Password: AnsiString; const PasswordHash: TSetupHash; const PasswordSalt: TSetupSalt): Boolean;
var
  Context: TSHA1Context;
  Hash: TSHA1Digest;
begin
  SHA1Init(Context);
  SHA1Update(Context, PAnsiChar('PasswordCheckHash')^, Length('PasswordCheckHash'));
  SHA1Update(Context, PasswordSalt, SizeOf(PasswordSalt));
  SHA1Update(Context, Pointer(Password)^, Length(Password)*SizeOf(Password[1]));
  Hash := SHA1Final(Context);
  Result := SHA1DigestsEqual(Hash, PasswordHash.SHA1);
end;

{$IFDEF UNICODE}
function TestPassword(Info: TMainInfo; const Password: string): Boolean;
begin
  Result := TestPassword(Info, FixPasswordEncoding(Info, Password));
end;
{$ENDIF}

function TestPassword(Info: TMainInfo; const Password: AnsiString): Boolean;
var
  Hash: TSetupHash;
begin
  Hash := Info.SetupHeader.PasswordHash;
  case Hash.HashType of
    htSHA256:
      Result := False; // not used for password checks (see TestPasswordIs64 instead), default to "failed"
    htSHA1:
      Result := TestPassword_SHA1(Password, Hash, Info.SetupHeader.PasswordSalt);
    htMD5:
      Result := TestPassword_MD5(Password, Hash, Info.SetupHeader.PasswordSalt);
  else
    Result := False;
  end;
end;

function TestPasswordIs64(Info: TMainInfo; const EncryptionKey: TSetupEncryptionKey): Boolean;
var
  Nonce : TSetupEncryptionNonce;
  Context : TChaCha20Context;
  PasswordTest : integer;
begin
  { Do same as compiler did in GeneratePasswordTest and compare results }
  Nonce := Info.SetupHeader.Is64Encryption.EncryptionBaseNonce;
  Nonce.RandomXorFirstSlice := Nonce.RandomXorFirstSlice xor -1;

  XChaCha20Init(Context, EncryptionKey[0], Length(EncryptionKey), Nonce, SizeOf(Nonce), 0);
  PasswordTest := 0;
  XChaCha20Crypt(Context, PasswordTest, PasswordTest, SizeOf(PasswordTest));

  Result := PasswordTest = Info.SetupHeader.Is64Encryption.PasswordTest;
end;

function GetSliceName(Info: TMainInfo; ASlice:integer):string;
var
  Major, Minor: Integer;
  Prefix:string;
begin
  if Info.Ver>=4107 then Prefix := PathChangeExt(PathExtractName(Info.SetupLdrOriginalFilename), '')
  else Prefix:=Info.SetupHeader.BaseFilename;
  Major := ASlice div Info.SetupHeader.SlicesPerDisk + 1;
  Minor := ASlice mod Info.SetupHeader.SlicesPerDisk;
  if Info.SetupHeader.SlicesPerDisk = 1 then
    Result := Format('%s-%d.bin', [Prefix, Major])
  else
    Result := Format('%s-%d%s.bin', [Prefix, Major, Chr(Ord('a') + Minor)]);
end;

function StrToOem(const s:string):AnsiString;
begin
  SetLength(Result, Length(s));
  CharToOemBuff(PChar(s),PAnsiChar(Result),length(s));
end;

function OemToStr(const s:AnsiString):string;
begin
  SetLength(Result, Length(s));
  OemToCharBuff(PAnsiChar(s),PChar(Result),length(s));
end;

// Delphi RTL does not check GetLastError() after getting GetFileType(hFile)==FILE_TYPE_UNKNOWN
// and incorrectly decides that an error occurred.
// This condition was observed when standard output was redirected by the parent process.
// Here is a workaround. Note that write/writeln now take exactly one argument of type string.
// It's also a good place to convert the characters into the OEM codepage.
// In case standard output is not a disk file, writes are split into 1K chunks to avoid
// out of memory errors.
procedure write(Info: TMainInfo; s:string);
var
  WriteOffset,BytesRemaining,BytesToWrite,BytesWritten:dword;
  ss: AnsiString;
begin
  ss:=StrToOem(s);
  WriteOffset:=0; BytesRemaining:=length(ss);
  while BytesRemaining>0 do begin
    if BytesRemaining>1024 then BytesToWrite:=1024 else BytesToWrite:=BytesRemaining;
    WriteFile(Info.StdOutputHandle,ss[WriteOffset+1],BytesToWrite,BytesWritten,nil);
{    if BytesWritten<>BytesToWrite then begin
      system.writeln('WriteFile failed. Requested=',BytestoWrite,' written=',BytesWritten,' error=',GetLastError);
      halt;
    end;}
    Inc(WriteOffset,BytesToWrite); Dec(BytesRemaining,BytesToWrite);
  end;
end;

procedure writeln(Info: TMainInfo; const s:string);
begin
  write(Info, s+#13#10);
end;

function AddFakeFile(Info: TMainInfo; const FileName: string; const FileContents:AnsiString; RenameNow:boolean=false):integer;
var
  pFileEntry:PSetupFileEntry;
  pFileLocationEntry:PSetupFileLocationEntry;
  SystemTime:TSystemTime;
  i:integer;
begin
  Result := -1;
  if (FileName = '') or (FileContents = '') then Exit;

  pFileLocationEntry:=AllocMem(sizeof(TSetupFileLocationEntry));
  with pFileLocationEntry^ do begin
    Contents:=FileContents;
    OriginalSize:=length(FileContents);
    DateTimeToSystemTime(Now,SystemTime);
    SystemTimeToFileTime(SystemTime,TimeStamp);
    PrimaryFileEntry:=-1;
  end;
  i:=Info.Entries[seFileLocation].Add(pFileLocationEntry);
  pFileEntry:=AllocMem(sizeof(TSetupFileEntry));
  with pFileEntry^ do begin
    FileType:=ftFakeFile;
    DestName:=FileName; // RenameFiles() will split this
    LocationEntry:=i;
    if RenameNow then begin
      SourceFilename:=FileName; DestName:=''; DestDir:='';
    end;
  end;

  Result:=Info.Entries[seFile].Add(pFileEntry);
end;

function FileContents(Filename:string):AnsiString;
var f:File;
begin
  AssignFile(f, Filename); Result:='';
  Reset(f,1);
  SetLength(Result, FileSize(f));
  BlockRead(f, Result[1], FileSize(f));
  CloseFile(f);
end;

function MakeDir(Dir: String): Boolean;
{ Returns True if a new directory was created }
var
  ErrorCode: DWORD;
begin
  Result := False;
  Dir := RemoveBackslash(ExpandFilename(Dir));
  if (Dir = '') or (AnsiLastChar(Dir)^ = ':') or (ExtractFilePath(Dir) = Dir) then
    Exit;
  if DirExists(Dir) then
    Exit;
  MakeDir(ExtractFilePath(Dir));
  if not CreateDirectory(PChar(Dir), nil) then begin
    ErrorCode := GetLastError;
    raise Exception.Create(FmtSetupMessage(msgLastErrorMessage,
      [FmtSetupMessage1(msgErrorCreatingDir, Dir), IntToStr(ErrorCode),
       SysErrorMessage(ErrorCode)]));
  end;
  Result := True;
end;

function GetCustomMessage(Info: TMainInfo; const AText: string): string;

  function ExtractFirstString(var AText: string): string;
  var
    SeparatorPos: integer;
  begin
    if AText <> '' then
    begin
      SeparatorPos := Pos(',', AText);
      if SeparatorPos >= 1 then
      begin
        Result := Copy(AText, 1, SeparatorPos - 1);
        Delete(AText, 1, SeparatorPos);
      end
      else
      begin
        Result := AText;
        AText := '';
      end;
    end
    else
      Result := '';
  end;

  function FindCustomMessage(const AName: string; out AText: string): boolean;
  var
    CustomMessageIndex: integer;
    CustomMessageEntry: PSetupCustomMessageEntry;
  begin
    Result := False;
    AText := '';
    for CustomMessageIndex := 0 to Info.Entries[seCustomMessage].Count - 1 do
    begin
      CustomMessageEntry := Info.Entries[seCustomMessage][CustomMessageIndex];
      if True
        and (CustomMessageEntry.LangIndex = 0) // main language
        and AnsiSameText(AName, CustomMessageEntry.Name)
      then
      begin
        AText := CustomMessageEntry.Value;
        Result := True;
        Break;
      end;
    end;
  end;

  function FormatCustomMessage(const AFormat: string; const AArgs: array of string): string;
  var
    Index: integer;
    ArgPos: integer;
    ArgValueIndex: integer;
    ArgValue: string;
  begin
    Result := AFormat;
    Index := 1;
    while Index < Length(Result) do
    begin
      case Result[Index] of
        '%':
          begin
            ArgPos := Index;
            Inc(Index);
            ArgValueIndex := 0;
            while (Index < Length(Result)) do
            begin
              case Result[Index] of
                '0'..'9':
                  begin
                    ArgValueIndex := ArgValueIndex * 10 + Ord(Result[Index]) - Ord('0');
                    Inc(Index);
                  end;
                else
                  Break;
              end;
            end;
            if True
              and (Index > ArgPos)
              and (ArgValueIndex > 0)
              and (ArgValueIndex <= Length(AArgs))
            then
            begin
              ArgValue := AArgs[ArgValueIndex - 1];
              Delete(Result, ArgPos, Index - ArgPos);
              Insert(ArgValue, Result, Index);
              Index := Index + Length(ArgValue);
            end;
          end;
        else
          Inc(Index);
      end;
    end;
  end;

var
  CustomMessage: string;
  CustomMessageName: string;
  CustomMessageFormat: string;
  CustomMessageArgs: array of string;
  Index: integer;
begin
  Result := AText;
  if True
    and (AText <> '')
    and (AText[Length(AText)] = '}')
    and AnsiStartsText(AText, '{cm:')
  then
  begin
    CustomMessage := Copy(AText, 5, Length(AText) - 5);
    CustomMessageName := ExtractFirstString(CustomMessage);
    SetLength(CustomMessageArgs, 0);
    while CustomMessage <> '' do
    begin
      Index := Length(CustomMessageArgs);
      SetLength(CustomMessageArgs, Index + 1);
      CustomMessageArgs[Index] := ExtractFirstString(CustomMessage);
    end;
    if FindCustomMessage(CustomMessageName, CustomMessageFormat) then
    begin
      Result := FormatCustomMessage(CustomMessageFormat, CustomMessageArgs);
    end;
  end;
end;

function FixPasswordEncoding(Info: TMainInfo; const Password: string) : AnsiString;
var
  tmpWS: WideString;
  nWideBytesCount: integer;
begin
  Result := AnsiString(Password);
  if (Info.VerIsUnicode) and (Password <> '') then
  begin
    tmpWS := Password;

    nWideBytesCount := Length(Result) * SizeOf(WideChar);
    SetLength(Result, nWideBytesCount);
    CopyMemory(@Result[1], PAnsiChar(@tmpWS[1]), nWideBytesCount);
  end;
end;

procedure GenerateEncryptionKey(const Password: String; const Salt: TSetupKDFSalt; const Iterations: Integer; out Key: TSetupEncryptionKey);
var
  SaltBytes,
  KeyBytes : TBytes;
  SaltSize,KeyLength : integer;
begin
  SaltSize := SizeOf(Salt);
  SetLength(SaltBytes, SaltSize);
  Move(Salt[0], SaltBytes[0], SaltSize);
  KeyLength := SizeOf(Key);
  KeyBytes := PBKDF2SHA256(Password, SaltBytes, Iterations, KeyLength);
  Move(KeyBytes[0], Key[0], KeyLength);
end;

end.