AWSTemplateFormatVersion: "2010-09-09"
Description: >- 
  This template creates an API Gateway and Lambda proxy to enable customized 
  HTTP headers for AS2 messages sent through AWS Transfer Family. It allows 
  you to dynamically modify the Subject header and add custom HTTP headers 
  required by your trading partners' AS2 servers. After deployment, update your 
  existing AWS Transfer Connector to use the API Gateway endpoint created by this
  stack. Supports both HTTP and HTTPS protocols. 
  WARNING: Only modify the Subject header or add headers your partner explicitly
  expects, as changing other headers may cause transfer failures.

Parameters:
  Environment:
    Type: String
    Default: prod
    AllowedValues:
      - dev
      - prod
    Description: Environment name
  PartnerHost:
    Type: String
    Description: IP address or hostname of the AS2 partner server
    Default: "partner-as2-server.example.com"
  PartnerPort:
    Type: String
    Description: Port number of the AS2 partner server (if not specified, defaults to 80 for HTTP and 443 for HTTPS)
  PartnerPath:
    Type: String
    Description: Path to the AS2 endpoint on the partner server
    Default: "/"
  ProtocolType:
    Type: String
    Description: Protocol to use for AS2 communication
    Default: HTTP
    AllowedValues:
      - HTTP
      - HTTPS

Conditions:
  IsHttps: !Equals [!Ref ProtocolType, "HTTPS"]
  IsPortSpecified: !Not [!Equals [!Ref PartnerPort, ""]]

Mappings:
  ProtocolSettings:
    HTTP:
      DefaultPort: 80
    HTTPS:
      DefaultPort: 443

Resources:
  # Lambda Layer for CA Certificates - Only created for HTTPS
  CertificatesLayer:
    Type: AWS::Lambda::LayerVersion
    Condition: IsHttps
    Properties:
      Description: Layer containing CA certificates for SSL validation
      Content:
        S3Bucket: !Ref "AWS::NoValue" # You need to upload zipped certificates to an S3 bucket
        S3Key: !Ref "AWS::NoValue" # and reference them here
      CompatibleRuntimes:
        - nodejs20.x

  # Lambda Function
  AS2HeadersFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.handler
      Runtime: nodejs20.x
      Layers: 
        !If [IsHttps, [!Ref CertificatesLayer], !Ref "AWS::NoValue"]
      Code:
        ZipFile: |
          // src/index.js
          const http = require('http');
          const https = require('https');
          const fs = require('fs');

          // Parameters from environment
          const partnerHost = process.env.PARTNER_HOST;
          const partnerPort = parseInt(process.env.PARTNER_PORT);
          const partnerPath = process.env.PARTNER_PATH;
          const protocolType = process.env.PROTOCOL_TYPE;

          const AS2_CONFIG = {
              host: partnerHost,
              port: partnerPort,
              path: partnerPath,
              method: 'POST'
          };

          exports.handler = async (event) => {
              console.log('Incoming event:', JSON.stringify(event, null, 2));
              console.log('Using protocol:', protocolType);

              try {
                  // Extract any dynamic values from query parameters or event
                  const filename = event.queryStringParameters?.filename || 'Unknown';
                  const timestamp = new Date().toISOString();

                  // Define your dynamic subject and other headers
                  const headers = {
                      ...event.headers,
                      'Subject': `File transfer: ${filename} - ${timestamp}`
                      // Add any other headers you need
                  };

                  console.log('Final headers:', JSON.stringify(headers, null, 2));

                  // Forward request to destination AS2 server
                  const response = await new Promise((resolve, reject) => {
                      const options = {
                          ...AS2_CONFIG,
                          headers: headers
                      };

                      console.log('Forwarding request with options:', JSON.stringify(options, null, 2));

                      // Use HTTPS or HTTP based on environment variable
                      const requestModule = protocolType === 'HTTPS' ? https : http;
                      
                      // Add HTTPS specific options if needed
                      if (protocolType === 'HTTPS') {
                          // You may need to load a custom CA certificate if the partner uses a non-public CA
                          // const caCert = fs.readFileSync('/opt/certificates/ca-cert.pem');
                          options.agent = new https.Agent({
                              // For custom CA: ca: caCert,
                              // For self-signed certs: rejectUnauthorized: false,
                          });
                      }

                      const req = requestModule.request(options, (res) => {
                          const chunks = [];

                          res.on('data', (chunk) => {
                              chunks.push(chunk);
                          });

                          res.on('end', () => {
                              const data = Buffer.concat(chunks);
                              resolve({
                                  statusCode: res.statusCode,
                                  data: data,
                                  headers: res.headers
                              });
                          });
                      });

                      // Add TLS event listeners for HTTPS
                      if (protocolType === 'HTTPS') {
                          req.on('socket', (socket) => {
                              socket.on('secureConnect', () => {
                                  console.log('TLS Handshake completed');
                                  console.log('Cipher:', socket.getCipher());
                                  console.log('Protocol:', socket.getProtocol());
                              });

                              socket.on('error', (error) => {
                                  console.error('Socket error:', error);
                              });
                          });
                      }

                      req.on('error', (error) => {
                          console.error('Error forwarding request:', error);
                          reject(error);
                      });

                      if (event.body) {
                          const body = event.isBase64Encoded
                              ? Buffer.from(event.body, 'base64')
                              : event.body;
                          req.write(body);
                      }

                      req.end();
                  });

                  console.log(response);
                  const responseObject = {
                      statusCode: response.statusCode,
                      body: response.data.toString('base64'),
                      headers: response.headers,
                      isBase64Encoded: true
                  };
                  console.log(responseObject);
                  return responseObject;

              } catch (error) {
                  console.error('Error occurred:', error);
                  console.error('Error stack:', error.stack);

                  return {
                      statusCode: 500,
                      body: JSON.stringify({
                          message: 'Error processing request',
                          error: error.message
                      })
                  };
              }
          };
      MemorySize: 128
      Timeout: 30
      Role: !GetAtt LambdaExecutionRole.Arn
      Environment:
        Variables:
          PARTNER_HOST: !Ref PartnerHost
          PARTNER_PORT: 
            !If [
              IsPortSpecified,
              !Ref PartnerPort,
              !If [
                IsHttps,
                !FindInMap [ProtocolSettings, HTTPS, DefaultPort],
                !FindInMap [ProtocolSettings, HTTP, DefaultPort]
              ]
            ]
          PARTNER_PATH: !Ref PartnerPath
          PROTOCOL_TYPE: !Ref ProtocolType

  # Lambda Execution Role
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  # HTTP API
  AS2HttpApi:
    Type: AWS::ApiGatewayV2::Api
    Properties:
      Name: !Sub ${AWS::StackName}-${Environment}-${AWS::Region}
      ProtocolType: HTTP
      Target: !GetAtt AS2HeadersFunction.Arn

  # API Integration
  AS2HttpApiIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref AS2HttpApi
      IntegrationType: AWS_PROXY
      IntegrationUri: !Sub arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${AS2HeadersFunction.Arn}/invocations
      PayloadFormatVersion: "2.0"

  # API Route
  AS2HttpApiRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref AS2HttpApi
      RouteKey: POST /process
      Target: !Sub integrations/${AS2HttpApiIntegration}

  # API Stage
  AS2HttpApiStage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      ApiId: !Ref AS2HttpApi
      StageName: !Ref Environment
      AutoDeploy: true

  # Lambda Permission for HTTP API
  LambdaHttpApiPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !Ref AS2HeadersFunction
      Principal: apigateway.amazonaws.com
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${AS2HttpApi}/*/*

Outputs:
  ApiEndpoint:
    Description: HTTPS API endpoint URL
    Value: !Sub https://${AS2HttpApi}.execute-api.${AWS::Region}.amazonaws.com/${Environment}/process

  LambdaArn:
    Description: Lambda Function ARN
    Value: !GetAtt AS2HeadersFunction.Arn

  LambdaRoleArn:
    Description: Lambda Execution Role ARN
    Value: !GetAtt LambdaExecutionRole.Arn

  ProtocolUsed:
    Description: Protocol used for AS2 communication
    Value: !Ref ProtocolType
