package de.rwth.idsg.bikeman.ixsi;

import de.rwth.idsg.bikeman.config.IxsiConfiguration;
import de.rwth.idsg.bikeman.ixsi.repository.SystemValidator;
import de.rwth.idsg.bikeman.web.rest.exception.DatabaseException;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import java.net.InetSocketAddress;
import java.util.Map;

/**
 * Created by max on 08/09/14.
 */
@Slf4j
@RequiredArgsConstructor
public class HandshakeInterceptor extends HttpSessionHandshakeInterceptor {
    private final SystemValidator systemValidator;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        InetSocketAddress inetSocketAddress = request.getRemoteAddress();
        log.trace("Starting handshake with {}", inetSocketAddress);

        String ip = inetSocketAddress.getAddress().getHostAddress();
        try {
            String systemId = systemValidator.getSystemID(ip);
            // to be be used in future
            attributes.put(IxsiConfiguration.SYSTEM_ID_KEY, systemId);
            return super.beforeHandshake(request, response, wsHandler, attributes);

        } catch (DatabaseException e) {
            log.error("Exception occurred", e);
            throw e;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception ex) {
        super.afterHandshake(request, response, wsHandler, ex);
        log.trace("Handshake complete with {}", request.getRemoteAddress());
    }
}

@@ -29,16 +29,55 @@
 .ixsi;%0A%0A
+import com.google.common.base.Strings;%0A
 import d
@@ -439,16 +439,158 @@
 sponse;%0A
+import org.springframework.web.context.request.RequestContextHolder;%0Aimport org.springframework.web.context.request.ServletRequestAttributes;%0A
 import o
@@ -739,30 +739,127 @@
 java
-.net.InetSocketAddress
+x.servlet.http.HttpServletRequest;%0Aimport java.net.InetSocketAddress;%0Aimport java.util.ArrayList;%0Aimport java.util.List
 ;%0Aim
@@ -1309,288 +1309,199 @@
     
-InetSocketAddress inetSocketAddress = request.getRemoteAddress();%0A        log.trace(%22Starting handshake with %7B%7D%22, inetSocketAddress);%0A%0A        String ip = inetSocketAddress.getAddress().getHostAddress();%0A        try %7B%0A            String systemId = systemValidator.getSystemID(ip);
+String systemId = getSystemID(request);%0A        if (systemId == null) %7B%0A            throw new DatabaseException(%22This ip address is not allowed for WebSocket communication%22);%0A        %7D else %7B
 %0A   
@@ -1694,118 +1694,8 @@
 es);
-%0A%0A        %7D catch (DatabaseException e) %7B%0A            log.error(%22Exception occurred%22, e);%0A            throw e;
 %0A   
@@ -2030,11 +2030,1437 @@
 );%0A    %7D
+%0A%0A    private String getSystemID(ServerHttpRequest request) %7B%0A        List%3CString%3E ipAddressList = getPossibleIpAddresses(request);%0A%0A        for (String ip : ipAddressList) %7B%0A            try %7B%0A                return systemValidator.getSystemID(ip);%0A            %7D catch (DatabaseException e) %7B%0A                // not in db, continue with the next in the list%0A            %7D%0A        %7D%0A        return null;%0A    %7D%0A%0A    private static List%3CString%3E getPossibleIpAddresses(ServerHttpRequest request) %7B%0A        List%3CString%3E ipAddressList = new ArrayList%3C%3E();%0A%0A        getFromProxy(request).stream()%0A                             .filter(s -%3E !Strings.isNullOrEmpty(s))%0A                             .forEach(ipAddressList::add);%0A%0A        ipAddressList.add(getFromRemote(request));%0A        ipAddressList.add(getFromContext());%0A%0A        return ipAddressList;%0A    %7D%0A%0A    private static List%3CString%3E getFromProxy(ServerHttpRequest request) %7B%0A        return request.getHeaders().get(%22X-FORWARDED-FOR%22);%0A    %7D%0A%0A    private static String getFromRemote(ServerHttpRequest request) %7B%0A        InetSocketAddress inetSocketAddress = request.getRemoteAddress();%0A        return inetSocketAddress.getAddress().getHostAddress();%0A    %7D%0A%0A    private static String getFromContext() %7B%0A        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest();%0A        return request.getRemoteAddr();%0A    %7D
 %0A%7D%0A

improve websocket ip checks during handshake