%% Copyright 2015 Kolab Systems AG (http://www.kolabsys.com)
%%
%% Aaron Seigo (Kolab Systems) <seigo a kolabsys.com>
%%
%% This program is free software: you can redistribute it and/or modify
%% it under the terms of the GNU General Public License as published by
%% the Free Software Foundation, either version 3 of the License, or
%% (at your option) any later version.
%%
%% This program is distributed in the hope that it will be useful,
%% but WITHOUT ANY WARRANTY; without even the implied warranty of
%% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%% GNU General Public License for more details.
%%
%% You should have received a copy of the GNU General Public License
%% along with this program.  If not, see <http://www.gnu.org/licenses/>.

-module(kolab_guam_session).

-behaviour(gen_server).

%% API
-export([ start_link/6 ]).

%% gen_server callbacks
-export([init/1, handle_call/3, handle_cast/2, handle_info/2, terminate/2, code_change/3]).

%% tests
-export([ process_client_data/3 ]).

-include("kolab_guam_session.hrl").
-include("compat.hrl").

%% public API
start_link(SupervisorPID, ListenSocket, ImapConfig, ImplicitTLS, TLSConfig, Rules) -> gen_server:start_link(?MODULE, [SupervisorPID, ListenSocket, ImapConfig, ImplicitTLS, TLSConfig, Rules], []).

%% gen_server API
init([SupervisorPID, ListenSocket, ServerConfig, ImplicitTLS, TLSConfig, Rules]) ->
    %% accepting a connection is blocking .. so do it async
    %% lager:debug("Starting a session handler on socket ~p for listener ~p", [ListenSocket, SupervisorPID]),
    process_flag(trap_exit, true),
    ActiveRules = init_rules(Rules),
    gen_server:cast(self(), accept),
    %% lager:debug("Rules are ~p from ~p", [ActiveRules, Rules]),
    { ok, #state{ listen_socket = ListenSocket, super_pid = SupervisorPID, client_implicit_tls = ImplicitTLS, tls_config = TLSConfig, server_config = ServerConfig, rules_deciding = ActiveRules } }.

handle_call(_Request, _From, State) ->
    { reply, ok, State }.

handle_cast(accept, State = #state{ server_config = ServerConfig } = State) ->
    %% try to rate limit our responses a bit here so that hammering the socket with connections is survivable
    timer:sleep(3),
    { ok, AcceptSocket, TLSActive } = accept_client(State),
    { ok, ImapSession } = eimap:start_link(ServerConfig),
    eimap:connect(ImapSession, self(), server_hello),
    { noreply, State#state{ listen_socket = undefined, socket = AcceptSocket, imap_session = ImapSession, client_tls_active = TLSActive } };
handle_cast(_Msg, State) ->
    { noreply, State }.

handle_info({ tcp_closed, _Socket }, State) ->
    %lager:debug("Client closed socket"),
    { stop, normal, State };
handle_info({ tcp_error, _Socket, _Error }, State) ->
    %lager:debug("Socket error"),
    { stop, normal, State };
handle_info({ ssl_closed, _Socket }, State) ->
    %lager:debug("Client closed socket"),
    { stop, normal, State };
handle_info({ ssl_error, _Socket, _Error }, State) ->
    %lager:debug("Socket error"),
    { stop, normal, State };
handle_info({ tcp, Socket, Data }, #state{ client_tls_active = false } = State) ->
    % lager:debug("Data coming in from client over TCP ~s", [Data]),
    {Acc, NewState} = process_client_data(Socket, Data, State),
    forward_literal_data(Acc, NewState),
    %Required because of socket upgrade, in which case we will end up with the ssl variant here.
    set_socket_active(NewState),
    { noreply, NewState };

handle_info({ ssl, Socket, Data }, State) ->
    % lager:debug("Data coming in from client over SSL, ~p", [Data]),
    {Acc, NewState} = process_client_data(Socket, Data, State),
    forward_literal_data(Acc, NewState),
    ssl:setopts(Socket, [{ active, once }]),
    { noreply, NewState };

handle_info({ server_hello, ServerHello }, #state{ imap_session = ImapSession, tls_config = TLSConfig, socket = Socket, client_implicit_tls = ImplicitTLS, client_tls_active = TLSActive, deflator = Deflator } = State) ->
    CorrectedHello = correct_hello(TLSActive, ImplicitTLS, TLSConfig, ServerHello),
    ServerIdent = proplists:get_value(server_id, ServerHello, <<>>),
    FullGreeting = <<"* OK [CAPABILITY ", CorrectedHello/binary, "] ", ServerIdent/binary, "\r\n">>,
    eimap:start_passthrough(ImapSession, self()),
    relay_response(Socket, postprocess_server_data(Deflator, FullGreeting), TLSActive),
    { noreply, State };
handle_info({ { rule_data, Module, ResponseToken }, Data }, #state{ rules_active = ActiveRules } = State) ->
    %lager:debug("Got back data requested by rule ~p: ~p", [Module, Data]),
    NewActiveRules =
    case proplists:get_value(Module, ActiveRules) of
        undefined -> ActiveRules;
        ModuleState ->
            NewModuleState = Module:imap_data(ResponseToken, Data, ModuleState),
            lists:reverse(lists:foldl(fun({ Rule, RuleState }, Acc) ->
                                              case Rule =:= Module of
                                                  true -> [{ Rule, NewModuleState }|Acc];
                                                  _ -> [{ Rule, RuleState }|Acc]
                                              end
                                      end,
                          [], ActiveRules))
    end,
    %TODO: should we also support non-active rules doing imapy things here?
    { noreply, State#state{ rules_active = NewActiveRules } };
handle_info({ imap_server_response, Data }, #state{ socket = Socket, imap_session = ImapSession, client_tls_active = TLS, deflator = Deflator, rules_active = ActiveRules } = State) ->
    lager:debug("FROM SERVER: ~p", [Data]),
    { ModifiedData, CurrentlyActiveRules } = apply_ruleset_serverside(ImapSession, Data, ActiveRules),
    relay_response(Socket, postprocess_server_data(Deflator, ModifiedData), TLS),
    NewSplitCommand = update_split_command_state(ModifiedData, State),
    { noreply, State#state{ rules_active = CurrentlyActiveRules, current_command_split = NewSplitCommand } };
handle_info({ 'EXIT', PID, _Reason }, #state { imap_session = PID } = State) ->
    { stop, normal, State#state{ imap_session = undefined } };
handle_info(Info, State) ->
    lager:debug("Received unexpected info... ~p", [Info]),
    { noreply, State }.

terminate(_Reason, #state{ inflator = Inflator, imap_session = ImapSession, deflator = Deflator, socket = Socket, client_implicit_tls = ImplicitTLS, client_tls_active = TLS }) ->
    close_zlib_handle(Inflator),
    close_zlib_handle(Deflator),
    close_socket(ImplicitTLS, TLS, Socket),
    case ImapSession of
        undefined -> ok;
        _ -> exit(ImapSession, kill)
    end,
    ok.

code_change(_OldVsn, State, _Extra) ->
    { ok, State }.

%% private API

%% update_split_command_state updates the split_command being held on to when we get a server response
%% in the case of "transactional" messages (such as authentication) where the client and server enter a bidirectional conversation
%% that is goverened by rules outside the the usual IMAP call/response pattern, we need to wait for the end of the server response
%% since this is relatively expensive due to having to scan the data for the tagged server response, and is not necessary for all other
%% IMAP commands, we shortcircuit when the command does not trigger a "transactional" interaction between client and server, and instead
%% just always reset the split data state at that point
update_split_command_state(Data, #state{ command_split_reset_trigger = reset_on_server_response, current_command_split = CurrentCommandSplit }) ->
    case CurrentCommandSplit of
        undefined -> undefined;
        { Tag, _Command, _Data } ->
            case binary:match(Data, <<Tag/binary, " ">>) of
                nomatch -> CurrentCommandSplit;
                { 0, _ } -> undefined;
                { Start, _ } ->
                    case binary:at(Data, Start - 1) of
                        $\n -> undefined;
                        _ -> CurrentCommandSplit
                    end
            end
    end;
update_split_command_state(_Data, _State) ->
    undefined.

update_split_command_state_client(ClientDataComponents, #state{ command_split_reset_trigger = reset_for_next_client_command, current_command_split = CurrentCommandSplit }) ->
    %If we have a currently active command (CurrentCommandSplit), then we reset only if we find a new command (and not just on any new line).
    %This is relevant for multi-line commands such as APPEND, which follow up with a continuation which we want to avoid buffering.
    case CurrentCommandSplit of
        undefined -> undefined;
        _ ->
            case ClientDataComponents of
                { <<>>, <<>>, <<>> } -> CurrentCommandSplit;
                { _Tag, <<>>, <<>> } -> CurrentCommandSplit;
                { <<>>, _Command, <<>> } -> CurrentCommandSplit;
                { _Tag, _Command, _Data } -> undefined
            end
    end;
update_split_command_state_client(_ClientDataComponents, _State) ->
    undefined.

set_socket_active(#state{ client_tls_active = true, socket = Socket}) -> ssl:setopts(Socket, [{ active, once }]);
set_socket_active(#state{ client_tls_active = false, socket = Socket}) -> inet:setopts(Socket, [{ active, once }]).


% Backwards compatibility
-ifdef(GUAM_OTP_21).
ssl_accept(Socket, TLSConfig) ->
    ssl:handshake(Socket, TLSConfig).
-else.
ssl_accept(Socket, TLSConfig) ->
    ssl:ssl_accept(Socket, TLSConfig).
-endif.

accept_client(#state{ client_implicit_tls = true, tls_config = TLSConfig, listen_socket = ListenSocket, super_pid = SupervisorPID }) ->
    AcceptSocket = accept_socket(ListenSocket, SupervisorPID),

    { ok, SSLSocket } =
    case ssl_accept(AcceptSocket, TLSConfig) of
        { ok, Socket } ->
            { ok, Socket };
        {error, {tls_alert, Reason}} ->
            lager:warning("TLS handshake failed with a tls_alert: ~p", [Reason]),
            ok = gen_tcp:close(AcceptSocket),
            exit(normal);
        {error, Reason} when Reason =:= timeout; Reason =:= closed ->
            lager:info("Socket closed or timed out during TLS handshake: ~p", [Reason]),
            ok = gen_tcp:close(AcceptSocket),
            exit(normal)
    end,

    %% prep for the next listen
    ok = ssl:setopts(SSLSocket, [{ active, once }, { mode, binary }]),
    % lager:info("~p All done!", [self()]),
    { ok, SSLSocket, true };
accept_client(#state{ listen_socket = ListenSocket, super_pid = SupervisorPID }) ->
    AcceptSocket = accept_socket(ListenSocket, SupervisorPID),
    ok = inet:setopts(AcceptSocket, [{ active, once }, { mode, binary }]),
    { ok, AcceptSocket, false }.

accept_socket(ListenSocket, SupervisorPID) ->
    AcceptResult = gen_tcp:accept(ListenSocket),
    %% start a new accepting process to replace this one, which is now in use
    supervisor:start_child(SupervisorPID, []),
    %% assert that the accept worked
    { ok, AcceptSocket } =
    case AcceptResult of
        { ok, Socket } ->
            { ok, Socket };
        {error, emfile} ->
            lager:error("Too many files open, you may be running into a file descriptor limitation."),
            exit(emfile)
    end,
    AcceptSocket.

close_zlib_handle(undefined) -> ok;
close_zlib_handle(Z) -> zlib:close(Z).

close_socket(_ImplicitTLS, _TLS, undefined) -> ok;
close_socket(_ImplicitTLS, true, Socket) -> ssl:close(Socket);
close_socket(true, _TLS, Socket) -> ssl:close(Socket);
close_socket(_ImplicitTLS, _TLS, Socket) -> gen_tcp:close(Socket).


forward_literal_data(<<>>, _State) ->
    {};
forward_literal_data(Data, #state{imap_session = ImapSession} = _State) ->
    eimap:passthrough_data(ImapSession, Data).

process_lines(_Socket, State, [], ContinuationBytes, Acc) ->
    {Acc, State#state{ continuation_bytes = ContinuationBytes }};
process_lines(Socket, State, [Line|MoreLines], 0, Acc) ->
    { _StrippedNextLine, NextContinuationBytes, _Result } = eimap_utils:num_literal_continuation_bytes(binary:part(Line, 0, byte_size(Line)-2)),
    {NewAcc, NewState} = process_client_line(Line, Socket, State, Acc),
    process_lines(Socket, NewState, MoreLines, NextContinuationBytes, NewAcc);
process_lines(Socket, State, [Line|MoreLines], ContinuationBytes, Acc) ->
    {LiteralPartOfLine, Remainder} = split_binary(Line, min(ContinuationBytes, size(Line))),
    Lines =
    case Remainder of
       <<>> -> MoreLines;
        _ -> [Remainder|MoreLines]
    end,
    process_lines(Socket, State, Lines, ContinuationBytes - size(LiteralPartOfLine), <<Acc/binary, LiteralPartOfLine/binary>>).

% Client data is processed as it comes in, but split up by lines if more than one line comes in at a time.
% The processing internally buffers data as necessary if it encounters incomplete lines.
process_client_data(Socket, Data, #state{ inflator = Inflator, continuation_bytes = ContinuationBytes } = State) ->
    % TODO: refactor so starttls and compress commands can be made into rules
    PreprocessData = preprocess_client_data(Inflator, Data, State),
    lager:debug("FROM CLIENT: ~p", [PreprocessData]),

    {Lines, NewLastPartialLine} = eimap_utils:split_lines(PreprocessData),
    {Acc, NewState} = process_lines(Socket, State#state{ buffered_client_data = <<>> }, Lines, ContinuationBytes, <<>>),
    #state{ buffered_client_data = DataToBuffer, continuation_bytes = NewContinuationBytes} = NewState,

    NewBuffer = <<DataToBuffer/binary, NewLastPartialLine/binary>>,
    %If we are inside a continuation, forward up to NewContinuationBytes
    {LiteralPartOfBuffer, NonLiteralPart} = split_binary(NewBuffer, min(NewContinuationBytes, size(NewBuffer))),

    {<<Acc/binary, LiteralPartOfBuffer/binary>>, NewState#state{ buffered_client_data = NonLiteralPart, continuation_bytes = NewContinuationBytes - size(LiteralPartOfBuffer)}}.


process_client_line(Data, Socket, #state{ rules_deciding = UndecidedRules, tls_config = TLSConfig, client_tls_active = TLS, rules_active = ActiveRules, socket = Socket, imap_session = ImapSession, deflator = Deflator, server_config = ServerConfig} = State, Acc) ->
    ClientDataComponents = eimap_utils:split_command_into_components(Data),
    case check_for_transmission_change_commands(TLS, TLSConfig, ClientDataComponents, Deflator, Socket) of
        { socket_upgraded, SSLSocket } ->
            %% if we have upgraded our socket, then do so to the backend if that hasn't happened auomatically
            case proplists:get_value(implicit_tls, ServerConfig, false) of
                false -> eimap:starttls(ImapSession, undefined, undefined);
                _ -> ok
            end,
            { Acc, State#state{ client_tls_active = true, socket = SSLSocket, buffered_client_data = <<>>, current_command_split = undefined, command_split_reset_trigger = undefined } };
        { compression, NewInflator, NewDeflator } ->
            eimap:compress(ImapSession), % TODO: make optional
            { Acc, State#state{ inflator = NewInflator, deflator = NewDeflator, buffered_client_data = <<>>, current_command_split = undefined, command_split_reset_trigger = undefined } };
        nochange ->
            %We  first check if we have to reset the split command before we process the new command in apply_ruleset_clientside
            NewSplit = update_split_command_state_client(ClientDataComponents, State),
            { ModifiedData, NewSplitCommand, NewSplitResetTrigger, NewUndecidedRules, NewActiveRules, PostAction } = apply_ruleset_clientside(ImapSession, Socket, Data, ClientDataComponents, NewSplit, UndecidedRules, ActiveRules),
            { SendThisData, BufferThisData } =
            case PostAction of
                perform_passthrough ->
                    {ModifiedData, <<>>};
                buffer_data ->
                    {<<>>, Data}
            end,
            { <<Acc/binary, SendThisData/binary>>, State#state{ rules_deciding = NewUndecidedRules, rules_active = NewActiveRules, buffered_client_data = BufferThisData, current_command_split = NewSplitCommand, command_split_reset_trigger = NewSplitResetTrigger } }
    end.

preprocess_client_data(undefined, Data, #state{ buffered_client_data = Buffered }) ->
    <<Buffered/binary, Data/binary>>;
preprocess_client_data(Z, Data, #state{ buffered_client_data = Buffered }) ->
    Inflated = iolist_to_binary(zlib:inflate(Z, Data)),
    <<Buffered/binary, Inflated/binary>>.

postprocess_server_data(undefined, Data) ->
    %% we aren't compressing so there is nothing to do
    Data;
postprocess_server_data(Z, Data) ->
    iolist_to_binary(zlib:deflate(Z, Data, sync)).

init_rules(RuleConfig) -> init_rule(RuleConfig, []).
init_rule([], Acc) -> Acc;
init_rule([{ RuleName, Config }|RuleConfig], Acc) ->
    Module = full_rule_name(RuleName),
    %% we try to new the module, but if something goes wrong, e.g. it does not exist,
    %% then we skip this config block because it is BROKEN
    try Module:new(Config) of
        ModuleState -> init_rule(RuleConfig, [{ Module, ModuleState }|Acc])
    catch
        Type:Error ->
            lager:warning("Could not create rule for ~p due to failure: ~p ~p", [RuleName, Type, Error]),
            init_rule(RuleConfig, Acc)
    end;
init_rule([_|RuleConfig], Acc) ->
    init_rule(RuleConfig, Acc).

full_rule_name(Module) when is_atom(Module) -> list_to_atom("kolab_guam_rule_" ++ atom_to_list(Module)).

apply_ruleset_serverside(ImapSession, ServerData, CurrentlyActiveRules) ->
    %TODO: allow undecided rules to opt-in here as well
    apply_next_rule_serverside(ImapSession, ServerData, [], CurrentlyActiveRules).

apply_next_rule_serverside(_ImapSession, ServerData, ActiveRulesAcc, []) -> { ServerData, lists:reverse(ActiveRulesAcc) };
apply_next_rule_serverside(ImapSession, ServerData, ActiveRulesAcc, [{ Module, RuleState } | ActiveRules]) ->
    %TODO: allow rules to remove themselves from the action during serverside processing?
    { ModifiedData, ModifiedRuleState } = Module:apply_to_server_message(ImapSession, ServerData, RuleState),
    apply_next_rule_serverside(ImapSession, ModifiedData, [{ Module, ModifiedRuleState } | ActiveRulesAcc], ActiveRules).

apply_ruleset_clientside(_ImapSession, _Socket, ClientData, _ClientDataComponents, _CurrentCommandSplit, [], []) ->
    { ClientData, undefined, [], [], [], perform_passthrough };
apply_ruleset_clientside(ImapSession, Socket, ClientData, ClientDataComponents, CurrentCommandSplit, UndecidedRules, CurrentlyActiveRules) ->
    { PostAction, SplitCommand, SplitResetTrigger } =
        case CurrentCommandSplit of
            undefined ->
                case ClientDataComponents of
                    % In case of an empty line ("\r\n") the line is complete (doesn't need to be buffered), but there is also no command or tag.
                    { <<>>, <<>>, <<>> } -> { perform_passthrough, CurrentCommandSplit, reset_for_next_client_command };
                    { _Tag, <<>>, <<>> } -> { buffer_data, undefined, reset_for_next_client_command };
                    { _Tag, Command, _Data } = Split -> { perform_passthrough, Split, when_to_reset_split(Command) }
                end;
            _ ->  { perform_passthrough, CurrentCommandSplit, reset_for_next_client_command }
        end,
    { StillUndecided, NewlyActive } = check_undecided(Socket, ClientData, SplitCommand, UndecidedRules),
    ActiveRules = CurrentlyActiveRules ++ NewlyActive,
    %lager:info("Active Rules: ~p", [ActiveRules]),
    { ModifiedData, ActiveRulesRun } = apply_next_rule_clientside(ImapSession, ClientData, SplitCommand, [], ActiveRules),
    { ModifiedData, SplitCommand, SplitResetTrigger, StillUndecided, ActiveRulesRun, PostAction }.

when_to_reset_split(<<"AUTHENTICATE">>) -> reset_on_server_response;
when_to_reset_split(<<"authenticate">>) -> reset_on_server_response;
when_to_reset_split(_) -> reset_for_next_client_command.

check_undecided(_Socket, _ClientData, undefined, Rules) ->
    %% if we do not have a properly split command ... do nothing!
    { Rules, [] };
check_undecided(Socket, ClientData, SplitCommand, Rules) -> check_next_undecided_rule(Socket, ClientData, SplitCommand, Rules, { [], [] }).
check_next_undecided_rule(_Socket, _ClientData, _SplitCommand, [], Accs) -> Accs;
check_next_undecided_rule(Socket, ClientData, SplitCommand, [Rule|Rules], { UndecidedAcc, NewActiveAcc }) ->
    { Module, RuleState } = Rule,
    %%lager:debug("Does ~p apply with state ~p? let's find out!", [Module, RuleState]),
    Application = Module:applies(Socket, ClientData, SplitCommand, RuleState),
    check_next_undecided_rule(Socket, ClientData, SplitCommand, Rules, applies(Module, Application, UndecidedAcc, NewActiveAcc)).

applies(Module, { true, RuleState }, UndecidedAcc, NewActiveAcc) -> { UndecidedAcc, [{ Module, RuleState }|NewActiveAcc] };
applies(_Module, { false, _RuleState }, UndecidedAcc, NewActiveAcc) -> { UndecidedAcc, NewActiveAcc };
applies(Module, { notyet, RuleState }, UndecidedAcc, NewActiveAcc) -> { [{ Module, RuleState }|UndecidedAcc], NewActiveAcc }.

apply_next_rule_clientside(_ImapSession, ClientData, _SplitCommand, ActiveRulesAcc, []) -> { ClientData, lists:reverse(ActiveRulesAcc) };
apply_next_rule_clientside(ImapSession, ClientData, SplitCommand, ActiveRulesAcc, [{ Module, RuleState }|Rules]) ->
    { Data, NewState } = Module:apply_to_client_message(ImapSession, ClientData, SplitCommand, RuleState),
    apply_next_rule_clientside(ImapSession, Data, SplitCommand, [{ Module, NewState } | ActiveRulesAcc], Rules).

relay_response(Socket, Data, false) ->
    %lager:debug("Sending over non-secure socket ..."),
    gen_tcp:send(Socket, Data);
relay_response(Socket, Data, _TLS) ->
    %lager:debug("Sending over TLS!"),
    ssl:send(Socket, Data).

check_for_transmission_change_commands(TLS, TLSConfig, ClientDataComponents, Deflator, Socket) ->
    {Tag, Command, _Data } = ClientDataComponents,
    case check_tls_state(TLS, TLSConfig, Command, Deflator, Socket, Tag) of
        nochange -> check_compress_request(Deflator, Command, Socket, TLS, Tag);
        Response -> Response
    end.

check_tls_state(false, TLSConfig, <<"STARTTLS">>, Deflator, Socket, Tag) -> start_client_tls(TLSConfig, Deflator, Socket, Tag);
check_tls_state(false, TLSConfig, <<"starttls">>, Deflator, Socket, Tag) -> start_client_tls(TLSConfig, Deflator, Socket, Tag);
check_tls_state(_TLS, _TLSConfig, _Buffer, _Deflator, _Socket, _Tag) -> nochange.

start_client_tls(TLSConfig, Deflator, Socket, Tag) ->
    Response = <<Tag/binary, " OK Begin TLS negotiation now\r\n">>,
    relay_response(Socket, postprocess_server_data(Deflator, Response), false),
    inet:setopts(Socket, [{ active, false }]), %% must be set to active false, otherwise can fail depending on timing
    { ok, SSLSocket } = ssl_accept(Socket, TLSConfig),
    { socket_upgraded, SSLSocket }.

check_compress_request(undefined, <<"COMPRESS">>, Socket, TLS, Tag) -> start_client_compression(Socket, TLS, Tag);
check_compress_request(undefined, <<"compress">>, Socket, TLS, Tag) -> start_client_compression(Socket, TLS, Tag);
check_compress_request(_Deflator, _Command, _Socket, _TLS, _Tag) -> nochange.

start_client_compression(Socket, TLS, Tag) ->
    Response = <<Tag/binary, " OK DEFLATE active\r\n">>,
    relay_response(Socket, postprocess_server_data(undefined, Response), TLS),
    %% create an inflate/deflate pair for use with the client
    Inflator = zlib:open(),
    ok = zlib:inflateInit(Inflator, -15),
    Deflator = zlib:open(),
    ok = zlib:deflateInit(Deflator, 1, deflated, -15, 8, default),
    { compression, Inflator, Deflator }.


-spec correct_hello(TLSActive :: true | false, ImplicitTLS :: true | false, TlSConfig :: [] | list(), ServerHello :: binary()) -> CorrectedHello :: binary().
correct_hello(true, true, _TLSConfig, ServerResponse) ->
    % the connection is already secured, so don't advertise starttls to the client
    ensure_hello_does_not_have_starttls(ServerResponse);
correct_hello(true, _ImplicitTLS, _TLSConfig, ServerResponse) ->
    % the connection is already secured, so don't advertise starttls to the client
    ensure_hello_does_not_have_starttls(ServerResponse);
correct_hello(_TLSActive, _ImplicitTLS, [], ServerResponse) ->
    % guam does not have a TLS config and so can not provide TLS to the client
    ensure_hello_does_not_have_starttls(ServerResponse);
correct_hello(_TLSActive, _ImplicitTLS, _TLSConfig, ServerResponse) ->
    % guam has a TLS config, and it is not currently active, so make sure to include
    % STARTTLS in our response regardless of what the backend says
    ensure_hello_has_starttls(ServerResponse).

ensure_hello_has_starttls(ServerResponse) ->
    ServerHello = proplists:get_value(capabilities, ServerResponse, <<>>),
    case binary:match(ServerHello, <<"STARTTLS">>) of
        nomatch -> add_starttls_to_capabilities(ServerHello);
        _ -> ServerHello
    end.

add_starttls_to_capabilities(ServerHello) ->
    case binary:match(ServerHello, <<"CAPABILITY ">>) of
        nomatch -> add_starttls_after_imap4_atom(ServerHello);
        { Start, End } ->
            Prefix = binary:part(ServerHello, 0, Start + End),
            Suffix = binary:part(ServerHello, Start + End, size(ServerHello) - Start - End),
            CorrectHello = <<Prefix/binary, "STARTTLS ", Suffix/binary>>,
            remove_auth_offers(CorrectHello)
    end.

add_starttls_after_imap4_atom(ServerHello) ->
    case binary:match(ServerHello, <<"IMAP4rev1 ">>) of
        nomatch -> <<"STARTTLS ", ServerHello/binary>>;
        { Start, End } ->
            Prefix = binary:part(ServerHello, 0, Start + End),
            Suffix = binary:part(ServerHello, Start + End, size(ServerHello) - Start - End),
            CorrectHello = <<Prefix/binary, "STARTTLS ", Suffix/binary>>,
            remove_auth_offers(CorrectHello)
    end.

ensure_hello_does_not_have_starttls(ServerResponse) ->
    ServerHello = proplists:get_value(capabilities, ServerResponse, <<>>),
    case binary:match(ServerHello, <<"STARTTLS">>) of
        nomatch -> ServerHello;
        { Start, End } ->
            Prefix = binary:part(ServerHello, 0, Start),
            Suffix = binary:part(ServerHello, Start + End, size(ServerHello) - Start - End),
            <<Prefix/binary, Suffix/binary>>
    end.

remove_auth_offers(ServerHello) ->
    case binary:match(ServerHello, <<"AUTH=">>) of
        nomatch -> ensure_advertise_login_disabled(ServerHello);
        { Start, _End } ->
            Prefix = binary:part(ServerHello, 0, Start),
            Suffix =
            case binary:match(ServerHello, <<" ">>, [{ scope, { Start, size(ServerHello) - Start } }]) of
                nomatch ->
                    %% end of the line, so no suffix
                    <<>>;
                { SpaceStart, SpaceEnd } ->
                    binary:part(ServerHello, SpaceStart + SpaceEnd, size(ServerHello) - SpaceStart - SpaceEnd)
            end,
            remove_auth_offers(<<Prefix/binary, Suffix/binary>>)
    end.

ensure_advertise_login_disabled(ServerHello) ->
    case binary:match(ServerHello, <<"LOGINDISABLED">>) of
        nomatch -> <<ServerHello/binary, " LOGINDISABLED">>;
        _ -> ServerHello
    end.

