package io.getlime.security.powerauth.lib.webflow.authentication.interceptor;

import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.lang.NonNull;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

/**
 * Web Socket handshake interceptor resolves the client IP address during Web Socket handshake and stores it into
 * message header attributes.
 *
 * @author Roman Strobl, roman.strobl@wultra.com
 */
public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {

    private final boolean forceIpv4;

    public WebSocketHandshakeInterceptor(boolean forceIpv4) {
        this.forceIpv4 = forceIpv4;
    }

    @Override
    public boolean beforeHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response,
                                   @NonNull WebSocketHandler wsHandler, @NonNull Map<String, Object> attributes) {
        // Set client_ip attribute in WebSocket session, either from the X-FORWARDED-FOR HTTP header, if it is not
        // available, use servlet request remote IP address.
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
            String ipAddress = servletRequest.getServletRequest().getHeader("X-FORWARDED-FOR");
            if (forceIpv4) {
                // IPv4 logic
                if (ipAddress == null || !isIpv4Address(ipAddress)) {
                    ipAddress = servletRequest.getServletRequest().getRemoteAddr();
                    if (!isIpv4Address(ipAddress)) {
                        // IP address is null in case IPv4 address could not be determined, it should not be sent to AFS
                        ipAddress = null;
                    }
                }
            } else if (ipAddress == null) {
                // IPv4 or IPv6 logic
                ipAddress = servletRequest.getServletRequest().getRemoteAddr();
            }

            attributes.put("client_ip", ipAddress);
        }
        return true;
    }

    @Override
    public void afterHandshake(@NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response,
                               @NonNull WebSocketHandler wsHandler, Exception exception) {
    }

    /**
     * Determine whether IP address is an IPv4 address.
     * @param address IP address as String.
     * @return Whether IP address is an IPv4 address.
     */
    private boolean isIpv4Address(String address) {
        // Source: https://www.oreilly.com/library/view/regular-expressions-cookbook/9781449327453/ch08s16.html
        return address.matches("^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\\.){3}↵\n" +
                "(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$");
    }
}

@@ -68,24 +68,125 @@
 terceptor;%0A%0A
+import com.google.common.net.InetAddresses;%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;%0A
 import org.s
@@ -223,32 +223,32 @@
 verHttpRequest;%0A
-
 import org.sprin
@@ -516,24 +516,100 @@
 terceptor;%0A%0A
+import javax.servlet.http.HttpServletRequest;%0Aimport java.net.Inet4Address;%0A
 import java.
@@ -901,16 +901,119 @@
 ptor %7B%0A%0A
+    private static final Logger logger = LoggerFactory.getLogger(WebSocketHandshakeInterceptor.class);%0A
     priv
@@ -1632,33 +1632,27 @@
         
+Http
 Servlet
-ServerHttp
 Request 
@@ -1668,16 +1668,17 @@
 quest = 
+(
 (Servlet
@@ -1703,16 +1703,37 @@
  request
+).getServletRequest()
 ;%0A      
@@ -1772,36 +1772,16 @@
 Request.
-getServletRequest().
 getHeade
@@ -1978,36 +1978,16 @@
 Request.
-getServletRequest().
 getRemot
@@ -2169,16 +2169,92 @@
  to AFS%0A
+                        logger.warn(%22IPv4 address could not be detected.%22);%0A
         
@@ -2456,28 +2456,8 @@
 est.
-getServletRequest().
 getR
@@ -3025,263 +3025,158 @@
     
-// Source: https://www.oreilly.com/library/view/regular-expressions-cookbook/9781449327453/ch08s16.html%0A        return address.matches(%22%5E(?:(?:25%5B0-5%5D%7C2%5B0-4%5D%5B0-9%5D%7C%5B01%5D?%5B0-9%5D%5B0-9%5D?)%5C%5C.)%7B3%7D%E2%86%B5%5Cn%22 +%0A                %22(?:25%5B0-5%5D%7C2%5B0-4%5D%5B0-9%5D%7C%5B01%5D?%5B0-9%5D%5B0-9%5D?)$%22);
+try %7B%0A            return InetAddresses.forString(address) instanceof Inet4Address;%0A        %7D catch (Exception e) %7B%0A            return false;%0A        %7D
 %0A   

Update logic for IP address detection