/*
    *  Licensed to the Apache Software Foundation (ASF) under one
    *  or more contributor license agreements.  See the NOTICE file
    *  distributed with this work for additional information
    *  regarding copyright ownership.  The ASF licenses this file
    *  to you under the Apache License, Version 2.0 (the
    *  "License"); you may not use this file except in compliance
    *  with the License.  You may obtain a copy of the License at
    *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   *  Unless required by applicable law or agreed to in writing,
   *  software distributed under the License is distributed on an
   *  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   *  KIND, either express or implied.  See the License for the
   *  specific language governing permissions and limitations
   *  under the License.
   *
   */
  package org.apache.mina.proxy.handlers.http.digest;
  
  import java.io.UnsupportedEncodingException;
  import java.security.NoSuchAlgorithmException;
  import java.security.SecureRandom;
  import java.util.Arrays;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.StringTokenizer;
  
  import org.apache.mina.core.filterchain.IoFilter.NextFilter;
  import org.apache.mina.proxy.ProxyAuthException;
  import org.apache.mina.proxy.handlers.http.AbstractAuthLogicHandler;
  import org.apache.mina.proxy.handlers.http.HttpProxyConstants;
  import org.apache.mina.proxy.handlers.http.HttpProxyRequest;
  import org.apache.mina.proxy.handlers.http.HttpProxyResponse;
  import org.apache.mina.proxy.session.ProxyIoSession;
  import org.apache.mina.proxy.utils.StringUtilities;
  import org.apache.mina.util.Base64;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * HttpDigestAuthLogicHandler.java - HTTP Digest authentication mechanism logic handler. 
   * 
   * @author The Apache MINA Project (dev@mina.apache.org)
   * @since MINA 2.0.0-M3
   */
  public class HttpDigestAuthLogicHandler extends AbstractAuthLogicHandler {
  
      private final static Logger logger = LoggerFactory
              .getLogger(HttpDigestAuthLogicHandler.class);
  
      /**
       * The challenge directives provided by the server.
       */
      private HashMap<String, String> directives = null;
  
      /**
       * The response received to the last request.
       */
      private HttpProxyResponse response;
  
      private static SecureRandom rnd;
  
      static {
          // Initialize secure random generator 
          try {
              rnd = SecureRandom.getInstance("SHA1PRNG");
          } catch (NoSuchAlgorithmException e) {
              throw new RuntimeException(e);
          }
      }
  
      public HttpDigestAuthLogicHandler(final ProxyIoSession proxyIoSession)
              throws ProxyAuthException {
          super(proxyIoSession);
  
          ((HttpProxyRequest) request).checkRequiredProperties(
                  HttpProxyConstants.USER_PROPERTY,
                  HttpProxyConstants.PWD_PROPERTY);
      }
  
      @Override
      public void doHandshake(NextFilter nextFilter) throws ProxyAuthException {
          logger.debug(" doHandshake()");
  
          if (step > 0 && directives == null) {
              throw new ProxyAuthException(
                      "Authentication challenge not received");
          }
          
          HttpProxyRequest req = (HttpProxyRequest) request;
          Map<String, List<String>> headers = req.getHeaders() != null ? req
                  .getHeaders() : new HashMap<String, List<String>>();
  
          if (step > 0) {
              logger.debug("  sending DIGEST challenge response");
  
             // Build a challenge response
             HashMap<String, String> map = new HashMap<String, String>();
             map.put("username", req.getProperties().get(
                     HttpProxyConstants.USER_PROPERTY));
             StringUtilities.copyDirective(directives, map, "realm");
             StringUtilities.copyDirective(directives, map, "uri");
             StringUtilities.copyDirective(directives, map, "opaque");
             StringUtilities.copyDirective(directives, map, "nonce");
             String algorithm = StringUtilities.copyDirective(directives,
                     map, "algorithm");
 
             // Check for a supported algorithm
             if (algorithm != null && !"md5".equalsIgnoreCase(algorithm)
                     && !"md5-sess".equalsIgnoreCase(algorithm)) {
                 throw new ProxyAuthException(
                         "Unknown algorithm required by server");
             }
 
             // Check for a supported qop
             String qop = directives.get("qop");
             if (qop != null) {
                 StringTokenizer st = new StringTokenizer(qop, ",");
                 String token = null;
 
                 while (st.hasMoreTokens()) {
                     String tk = st.nextToken();
                     if ("auth".equalsIgnoreCase(token)) {
                         break;
                     }
 
                     int pos = Arrays.binarySearch(
                             DigestUtilities.SUPPORTED_QOPS, tk);
                     if (pos > -1) {
                         token = tk;
                     }
                 }
 
                 if (token != null) {
                     map.put("qop", token);
 
                     byte[] nonce = new byte[8];
                     rnd.nextBytes(nonce);
 
                     try {
                         String cnonce = new String(Base64
                                 .encodeBase64(nonce), proxyIoSession
                                 .getCharsetName());
                         map.put("cnonce", cnonce);
                     } catch (UnsupportedEncodingException e) {
                         throw new ProxyAuthException(
                                 "Unable to encode cnonce", e);
                     }
                 } else {
                     throw new ProxyAuthException(
                             "No supported qop option available");
                 }
             }
 
             map.put("nc", "00000001");
             map.put("uri", req.getHttpURI());
 
             // Compute the response
             try {
                 map.put("response", DigestUtilities
                         .computeResponseValue(proxyIoSession.getSession(),
                                 map, req.getHttpVerb().toUpperCase(),
                                 req.getProperties().get(
                                         HttpProxyConstants.PWD_PROPERTY),
                                 proxyIoSession.getCharsetName(), response
                                         .getBody()));
 
             } catch (Exception e) {
                 throw new ProxyAuthException(
                         "Digest response computing failed", e);
             }
 
             // Prepare the challenge response header and add it to the 
             // request we will send
             StringBuilder sb = new StringBuilder("Digest ");
             boolean addSeparator = false;
 
             for (String key : map.keySet()) {
 
                 if (addSeparator) {
                     sb.append(", ");
                 } else {
                     addSeparator = true;
                 }
 
                 boolean quotedValue = !"qop".equals(key)
                         && !"nc".equals(key);
                 sb.append(key);
                 if (quotedValue) {
                     sb.append("=\"").append(map.get(key)).append('\"');
                 } else {
                     sb.append('=').append(map.get(key));
                 }
             }
 
             StringUtilities.addValueToHeader(headers,
                     "Proxy-Authorization", sb.toString(), true);
         }
 
         addKeepAliveHeaders(headers);
         req.setHeaders(headers);
 
         writeRequest(nextFilter, req);
         step++;
     }
 
     @Override
     public void handleResponse(final HttpProxyResponse response)
             throws ProxyAuthException {
         this.response = response;
 
         if (step == 0) {
             if (response.getStatusCode() != 401
                     && response.getStatusCode() != 407) {
                 throw new ProxyAuthException(
                         "Received unexpected response code ("
                                 + response.getStatusLine() + ").");
             }
 
             // Header should look like this
             // Proxy-Authenticate: Digest still_some_more_stuff
             List<String> values = response.getHeaders().get(
                     "Proxy-Authenticate");
             String challengeResponse = null;
 
             for (String s : values) {
                 if (s.startsWith("Digest")) {
                     challengeResponse = s;
                     break;
                 }
             }
 
             if (challengeResponse == null) {
                 throw new ProxyAuthException(
                         "Server doesn't support digest authentication method !");
             }
 
             try {
                 directives = StringUtilities.parseDirectives(challengeResponse
                         .substring(7).getBytes(proxyIoSession.getCharsetName()));
             } catch (Exception e) {
                 throw new ProxyAuthException(
                         "Parsing of server digest directives failed", e);
             }
             step = 1;
         } else {
             throw new ProxyAuthException("Received unexpected response code ("
                     + response.getStatusLine() + ").");
         }
     }
 }