%%
%% %CopyrightBegin%
%%
%% Copyright Ericsson AB 2005-2012. All Rights Reserved.
%%
%% The contents of this file are subject to the Erlang Public License,
%% Version 1.1, (the "License"); you may not use this file except in
%% compliance with the License. You should have received a copy of the
%% Erlang Public License along with this software. If not, it can be
%% retrieved online at http://www.erlang.org/.
%%
%% Software distributed under the License is distributed on an "AS IS"
%% basis, WITHOUT WARRANTY OF ANY KIND, either express or implied. See
%% the License for the specific language governing rights and limitations
%% under the License.
%%
%% %CopyrightEnd%
%%

%%

%%% Description: SSH file handling

-module(ssh_file).

-include_lib("public_key/include/public_key.hrl").
-include_lib("kernel/include/file.hrl").

-include("ssh.hrl").

-export([public_host_dsa_key/2,private_host_dsa_key/2,
	 public_host_rsa_key/2,private_host_rsa_key/2,
	 %%public_host_key/2,
	 private_host_key/2,
	 lookup_host_key/3, add_host_key/3,
	 lookup_user_key/4, ssh_dir/2, file_name/3]).

-export([private_identity_key/2, 
	 public_identity_key/2]).

-export([encode_public_key/1, decode_public_key_v2/2]).

-define(PERM_700, 8#700).
-define(PERM_644, 8#644).


%% API
public_host_dsa_key(Type, Opts) ->
    File = file_name(Type, "ssh_host_dsa_key.pub", Opts),
    decode(File, public_key).

private_host_dsa_key(Type, Opts) ->
    File = file_name(Type, "ssh_host_dsa_key", Opts),
    Password = proplists:get_value(password, Opts, ignore),
    decode(File, Password).

public_host_rsa_key(Type, Opts) ->
    File = file_name(Type, "ssh_host_rsa_key.pub", Opts),
    decode(File, public_key).

private_host_rsa_key(Type, Opts) ->
    File = file_name(Type, "ssh_host_rsa_key", Opts),
    Password = proplists:get_value(password, Opts, ignore),
    decode(File, Password).

%% public_host_key(Type, Opts) ->
%%     File = file_name(Type, "ssh_host_key", Opts),
%%     decode(File, public_key).
 
private_host_key(Type, Opts) ->
    File = file_name(Type, "ssh_host_key", Opts),
    Password = proplists:get_value(password, Opts, ignore),
    decode(File, Password).

private_identity_key(Alg, Opts) ->
    File = file_name(user, identity_key_filename(Alg), Opts),
    Password = proplists:get_value(password, Opts, ignore),
    decode(File, Password).

public_identity_key(Alg, Opts) ->
    File = file_name(user, identity_key_filename(Alg) ++ ".pub", Opts),
    decode(File, public_key).

encode_public_key(#'RSAPrivateKey'{publicExponent = E, modulus = N}) ->
    ssh_bits:encode(["ssh-rsa",E,N], [string,mpint,mpint]);
encode_public_key(#'DSAPrivateKey'{p = P, q = Q, g = G, y = Y}) ->
    ssh_bits:encode(["ssh-dss",P,Q,G,Y], [string,mpint,mpint,mpint,mpint]).

decode(File, Password) ->
    try 
	{ok, decode_ssh_file(read_ssh_file(File), Password)}
    catch 
	throw:Reason ->
	    {error, Reason};
	error:Reason ->
	    {error, Reason}
    end.      

read_ssh_file(File) ->
    {ok, Bin} = file:read_file(File),
    Bin.

%% Public key
decode_ssh_file(SshBin, public_key) ->
    public_key:ssh_decode(SshBin, public_key);

%% Private Key
decode_ssh_file(Pem, Password) ->
    case public_key:pem_decode(Pem) of
	[{_, _, not_encrypted} = Entry]  -> 
	    public_key:pem_entry_decode(Entry);
	[Entry] when Password =/= ignore ->
	    public_key:pem_entry_decode(Entry, Password);
	 _ ->
	    throw("No pass phrase provided for private key file")
    end.
    

%% lookup_host_key
%% return {ok, Key(s)} or {error, not_found}
%%

lookup_host_key(Host, Alg, Opts) ->
    Host1 = replace_localhost(Host),
    do_lookup_host_key(Host1, Alg, Opts).
	    

add_host_key(Host, Key, Opts) ->
    Host1 = add_ip(replace_localhost(Host)),
    KnownHosts = file_name(user, "known_hosts", Opts),
    case file:open(KnownHosts, [write,append]) of
   	{ok, Fd} ->
	    ok = file:change_mode(KnownHosts, ?PERM_644),
   	    Res = add_key_fd(Fd, Host1, Key),
   	    file:close(Fd),
   	    Res;
   	Error ->
   	    Error
    end.

lookup_user_key(Key, User, Alg, Opts) ->
    SshDir = ssh_dir({remoteuser,User}, Opts),
    case lookup_user_key_f(Key, User, SshDir, Alg, "authorized_keys", Opts) of
	{ok, Key} ->
	    {ok, Key};
	_ ->
	    lookup_user_key_f(Key, User, SshDir, Alg,  "authorized_keys2", Opts)
    end.


%%
%% Utils
%%

%% server use this to find individual keys for
%% an individual user when user tries to login
%% with publickey
ssh_dir({remoteuser, User}, Opts) ->
    case proplists:get_value(user_dir_fun, Opts) of
	undefined ->
	    case proplists:get_value(user_dir, Opts) of
		undefined ->
		    default_user_dir();
		Dir ->
		    Dir
	    end;
	FUN ->
	    FUN(User)
    end;

%% client use this to find client ssh keys
ssh_dir(user, Opts) ->
    case proplists:get_value(user_dir, Opts, false) of
	false -> default_user_dir();
	D -> D
    end;

%% server use this to find server host keys
ssh_dir(system, Opts) ->
    proplists:get_value(system_dir, Opts, "/etc/ssh").


file_name(Type, Name, Opts) ->
    FN = filename:join(ssh_dir(Type, Opts), Name),
    FN.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% in: "host" out: "host,1.2.3.4.
add_ip(Host)                                                             ->
    case inet:getaddr(Host, inet) of
	{ok, Addr} ->
	    case ssh_connection:encode_ip(Addr) of
		false -> Host;
		IPString -> Host ++ "," ++ IPString
	    end;
	_ -> Host
    end.    

replace_localhost("localhost") ->
    {ok, Hostname} = inet:gethostname(),
    Hostname;
replace_localhost(Host) ->
    Host.

do_lookup_host_key(Host, Alg, Opts) ->
    case file:open(file_name(user, "known_hosts", Opts), [read, binary]) of
	{ok, Fd} ->
	    Res = lookup_host_key_fd(Fd, Host, Alg),
	    file:close(Fd),
	    {ok, Res};
	{error, enoent} -> {error, not_found};
	Error -> Error
    end.

identity_key_filename("ssh-dss") -> "id_dsa";
identity_key_filename("ssh-rsa") -> "id_rsa".

decode_public_key_v2(K_S, "ssh-rsa") ->
    case ssh_bits:decode(K_S,[string,mpint,mpint]) of
	["ssh-rsa", E, N] ->
	    {ok, #'RSAPublicKey'{publicExponent = E, modulus = N}};
	_ ->
	    {error, bad_format}
    end;
decode_public_key_v2(K_S, "ssh-dss") ->  
     case ssh_bits:decode(K_S,[string,mpint,mpint,mpint,mpint]) of
     	["ssh-dss",P,Q,G,Y] ->
     	    {ok, {Y, #'Dss-Parms'{p = P, q = Q, g = G}}};
     	_A ->
     	    {error, bad_format}
     end;
decode_public_key_v2(_, _) ->
    {error, bad_format}.
    
lookup_host_key_fd(Fd, Host, KeyType) ->
    case io:get_line(Fd, '') of
	eof ->
	    {error, not_found};
	Line ->
	    case public_key:ssh_decode(Line, known_hosts) of
		[{Key, Attributes}] ->
		    handle_host(Fd, Host, proplists:get_value(hostnames, Attributes), Key, KeyType);
		[] ->
		    lookup_host_key_fd(Fd, Host, KeyType)
	    end
    end.

handle_host(Fd, Host, HostList, Key, KeyType) ->
    Host1 = host_name(Host),
    case lists:member(Host1, HostList) and key_match(Key, KeyType) of
	true ->
	    Key;
	false ->
	    lookup_host_key_fd(Fd, Host, KeyType)
    end.

host_name(Atom) when is_atom(Atom) ->
    atom_to_list(Atom);
host_name(List) ->
    List.

key_match(#'RSAPublicKey'{}, "ssh-rsa") ->
    true;
key_match({_, #'Dss-Parms'{}}, "ssh-dss") ->
    true;
key_match(_, _) ->
    false.

add_key_fd(Fd, Host,Key) ->
    SshBin = public_key:ssh_encode([{Key, [{hostnames, [Host]}]}], known_hosts),
    file:write(Fd, SshBin).

lookup_user_key_f(_, _User, [], _Alg, _F, _Opts) ->
    {error, nouserdir};
lookup_user_key_f(_, _User, nouserdir, _Alg, _F, _Opts) ->
    {error, nouserdir};
lookup_user_key_f(Key, _User, Dir, _Alg, F, _Opts) ->
    FileName = filename:join(Dir, F),
    case file:open(FileName, [read, binary]) of
	{ok, Fd} ->
	    Res = lookup_user_key_fd(Fd, Key),
	    file:close(Fd),
	    Res;
	{error, Reason} ->
	    {error, {{openerr, Reason}, {file, FileName}}}
    end.

lookup_user_key_fd(Fd, Key) ->
    case io:get_line(Fd, '') of
	eof ->
	    {error, not_found};
	Line ->
	    case public_key:ssh_decode(Line, auth_keys) of
		[{AuthKey, _}] ->
		    case is_auth_key(Key, AuthKey) of
			true ->
			    {ok, Key};
			false ->
			    lookup_user_key_fd(Fd, Key)
		    end;
		[] ->
		    lookup_user_key_fd(Fd, Key)
	    end
    end.

is_auth_key(Key, Key) -> 
    true;
is_auth_key(_,_) -> 
    false.


default_user_dir()->
    {ok,[[Home|_]]} = init:get_argument(home),
    UserDir = filename:join(Home, ".ssh"),
    ok = filelib:ensure_dir(filename:join(UserDir, "dummy")),
    ok = file:change_mode(UserDir, ?PERM_700),
    UserDir.