1 files changed, 18 insertions, 16 deletions

diff --git a/id/server/idserverlib/src/main/java/at/gv/egovernment/moa/id/protocols/stork2/AttributeCollector.java b/id/server/idserverlib/src/main/java/at/gv/egovernment/moa/id/protocols/stork2/AttributeCollector.java
index bd32bfc78..5f46153af 100644
--- a/id/server/idserverlib/src/main/java/at/gv/egovernment/moa/id/protocols/stork2/AttributeCollector.java
+++ b/id/server/idserverlib/src/main/java/at/gv/egovernment/moa/id/protocols/stork2/AttributeCollector.java
@@ -19,8 +19,6 @@ import eu.stork.peps.auth.commons.IPersonalAttributeList;
 import eu.stork.peps.auth.commons.PEPSUtil;
 import eu.stork.peps.auth.commons.PersonalAttribute;
 import eu.stork.peps.auth.commons.PersonalAttributeList;
-import eu.stork.peps.auth.commons.STORKAuthnRequest;
-import eu.stork.peps.auth.commons.STORKAuthnResponse;
 import eu.stork.peps.auth.engine.STORKSAMLEngine;
 import eu.stork.peps.exceptions.STORKSAMLEngineException;
 import org.apache.velocity.Template;
@@ -183,15 +181,18 @@ public class AttributeCollector implements IAction {
      * @throws MOAIDException the mOAID exception
      */
     private void generateSTORKResponse(DataContainer container) throws MOAIDException {
-        STORKAuthnResponse authnResponse = container.getResponse();
-        STORKAuthnRequest authnRequest = container.getRequest();
+    	MOASTORKRequest request = container.getRequest();
+        MOASTORKResponse response = container.getResponse();
 
         try {
             //Get SAMLEngine instance
             STORKSAMLEngine engine = STORKSAMLEngine.getInstance("VIDP");
             Logger.debug("Starting generation of SAML response");
-            authnResponse = engine.generateSTORKAuthnResponse(authnRequest, authnResponse, container.getRemoteAddress(), false);
-
+			if(response.isAuthnResponse())
+				response.setSTORKAuthnResponse(engine.generateSTORKAuthnResponse(request.getStorkAuthnRequest(), response.getStorkAuthnResponse(), container.getRemoteAddress(), false));
+			else
+				response.setSTORKAttrResponse(engine.generateSTORKAttrQueryResponse(request.getStorkAttrQueryRequest(), response.getStorkAttrQueryResponse(), container.getRemoteAddress(), "", false));
+				
             //generateSAML Token
             Logger.info("SAML response succesfully generated!");
         } catch (STORKSAMLEngineException e) {
@@ -200,11 +201,6 @@ public class AttributeCollector implements IAction {
         }
 
         Logger.info("STORK SAML Response message succesfully generated ");
-        Logger.debug("authn saml plain:" + authnResponse.getTokenSaml());
-		Logger.debug("authn saml string:" + new String(authnResponse.getTokenSaml()));
-		Logger.debug("authn saml encodedx: " + PEPSUtil.encodeSAMLToken(authnResponse.getTokenSaml()));
-
-        container.setResponse(authnResponse);
     }
 
     /**
@@ -214,17 +210,23 @@ public class AttributeCollector implements IAction {
      * @param container the container
      */
     private void generateRedirectResponse(HttpServletResponse httpResp, DataContainer container) {
-        STORKAuthnResponse authnResponse = container.getResponse();
-        STORKAuthnRequest authnRequest = container.getRequest();
+        MOASTORKResponse authnResponse = container.getResponse();
+        MOASTORKRequest authnRequest = container.getRequest();
 
         // preparing redirection for the client
         try {
             VelocityEngine velocityEngine = VelocityProvider.getClassPathVelocityEngine();
             Template template = velocityEngine.getTemplate("/resources/templates/stork2_postbinding_template.html");
             VelocityContext context = new VelocityContext();
-
-            context.put("SAMLResponse", PEPSUtil.encodeSAMLToken(authnResponse.getTokenSaml()).getBytes());
-            Logger.debug("SAMLResponse original: " + new String(authnResponse.getTokenSaml()).getBytes());
+            
+            byte[] blob;
+			if(authnRequest.isAttrRequest())
+            	blob = authnResponse.getStorkAttrQueryResponse().getTokenSaml();
+            else
+            	blob = authnResponse.getStorkAuthnResponse().getTokenSaml();
+
+            context.put("SAMLResponse", PEPSUtil.encodeSAMLToken(blob).getBytes());
+            Logger.debug("SAMLResponse original: " + new String(blob).getBytes());
 
             Logger.debug("Putting assertion consumer url as action: " + authnRequest.getAssertionConsumerServiceURL());
             context.put("action", authnRequest.getAssertionConsumerServiceURL());