///////////////////////////////////////////////////////////////////////////////
//
// Copyright (c) 2000-2019 Ericsson Telecom AB
//
// All rights reserved. This program and the accompanying materials
// are made available under the terms of the Eclipse Public License v2.0
// which accompanies this distribution, and is available at
// https://www.eclipse.org/org/documents/epl-2.0/EPL-2.0.html
///////////////////////////////////////////////////////////////////////////////
//
//  File:               STUN_EncDec.cc
//  Description:        Encoder/Decoder functions for STUN Protocol module
//  Rev:                R2B
//  Prodnr:             CNL 113 778
//  Updated:            2013-07-30
//  Contact:            http://ttcn.ericsson.se
//  Reference:          
//
//
#include "STUN_Types.hh"
#include <openssl/hmac.h>
#include <openssl/rand.h>
#include <zlib.h>

const int MESSAGE_INTEGRITY_LENGTH = 24;
const int MESSAGE_INTEGRITY_VALUE_LENGTH = 20;
const int FINGERPRINT_LENGTH = 8;
const int FINGERPRINT_VALUE_LENGTH = 4;
const int STUN_MESSAGE_HEADER_LENGTH = 20;

using namespace General__Types;

namespace STUN__Types{
static const unsigned char EMPTY_OCTETS_16[] = {
  0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,	0, 0, 0, 0
};
																								 
static const unsigned char EMPTY_OCTETS_20[] = {
  0, 0, 0, 0,	0, 0, 0, 0,	0, 0, 0, 0,	0, 0, 0, 0,	0, 0, 0, 0
};

static const OCTETSTRING OCT16_EMPTY(16, EMPTY_OCTETS_16);
static const OCTETSTRING OCT20_EMPTY(20, EMPTY_OCTETS_20);


OCTETSTRING get_HMAC(const unsigned char* pl_stream, const int& pl_stream_length,
                      const unsigned char* pl_key, const int& pl_key_length) {
  if (pl_stream_length > 0) {

    unsigned char md[EVP_MAX_MD_SIZE];
    unsigned int md_len;

    HMAC(EVP_sha1(), pl_key, pl_key_length, pl_stream, pl_stream_length, md, &md_len);

    OCTETSTRING result(md_len, md);

    return result;
  }
  
  TTCN_warning("Message length too short");
  return OCT20_EMPTY;
} //get_HMAC


void insert_Fingerprint(TTCN_Buffer& pl_buffer) {
  const unsigned char xor_value[] = {'\x53', '\x54', '\x55', '\x4E'};
  unsigned char crc_stream[4];
  unsigned char attribute_head[4];
  //calculate the fingerprint attribute
  uLong crc = crc32(0L, pl_buffer.get_data(), pl_buffer.get_len());
  //insert the attribute type and length
  attribute_head[0] = (STUN__Attribute__Type::FINGERPRINT >> 8) & 0xFF;
  attribute_head[1] = (STUN__Attribute__Type::FINGERPRINT & 0xFF);
  attribute_head[2] = (4 >> 8) & 0xFF;
  attribute_head[3] = (4 & 0xFF);
  pl_buffer.put_s(4, attribute_head);
  //insert the attribute value
  crc_stream[0] = ((crc >> 24) & 0xFF) xor xor_value[0];
  crc_stream[1] = ((crc >> 16) & 0xFF) xor xor_value[1];
  crc_stream[2] = ((crc >> 8) & 0xFF)  xor xor_value[2];
  crc_stream[3] = (crc & 0xFF) xor xor_value[3];
  pl_buffer.put_s(4, crc_stream);
} //insert_Fingerprint


void insert_HMAC(TTCN_Buffer& pl_buffer, const unsigned char* pl_key, const int& pl_key_length){
  unsigned char attribute_head[4];
  attribute_head[0] = (STUN__Attribute__Type::MESSAGE__INTEGRITY >> 8) & 0xFF;
  attribute_head[1] = (STUN__Attribute__Type::MESSAGE__INTEGRITY & 0xFF);
  attribute_head[2] = (20 >> 8) & 0xFF;
  attribute_head[3] = (20 & 0xFF);

  //pl_stream does not contain the message integrity attribute at all(24 bytes)
  OCTETSTRING hmac = get_HMAC(pl_buffer.get_data(), pl_buffer.get_len(), pl_key, pl_key_length);

  pl_buffer.put_s(4, attribute_head);
  pl_buffer.put_os(hmac);

} //insert_HMAC

void change_Message_Length(TTCN_Buffer& pl_buffer, const unsigned int& new_message_length) {
   unsigned char* buffer_data = NULL;
   size_t ss = 1;
   pl_buffer.get_end(buffer_data, ss);

   *(buffer_data-pl_buffer.get_len()+2) = (new_message_length >> 8) & 0xFF;
   *(buffer_data-pl_buffer.get_len()+3) = (new_message_length & 0xFF);
}// change_Message_Length

BOOLEAN ef__STUN__Check__Message__Integrity(const OCTETSTRING& pl__stream,
                                            const OCT16& pl__key,
                                            const STUN__Attribute__Type& pl__last__attribute__type) {
  int offset = -1;
  switch(pl__last__attribute__type) {
    case STUN__Attribute__Type::MESSAGE__INTEGRITY:
      offset = pl__stream.lengthof() - MESSAGE_INTEGRITY_LENGTH;
    break;
    case STUN__Attribute__Type::FINGERPRINT:
      offset = pl__stream.lengthof() - FINGERPRINT_LENGTH - MESSAGE_INTEGRITY_LENGTH;
    break;
    default:
      TTCN_warning("Last attribute is not Message Integrity");
      return true;
  };
  
  if(offset < 0){
    TTCN_error("Negative offset.");
  }

  //offset+1 == attribute type
	if(((const unsigned char*)pl__stream)[offset+1] == STUN__Attribute__Type::MESSAGE__INTEGRITY) {
	  OCTETSTRING stream_MAC = substr(pl__stream, offset+4, MESSAGE_INTEGRITY_VALUE_LENGTH);
	  TTCN_Buffer buffer(substr(pl__stream, 0, offset));
	  change_Message_Length(buffer, (offset+MESSAGE_INTEGRITY_LENGTH-STUN_MESSAGE_HEADER_LENGTH));
	  OCTETSTRING calculated_MAC = get_HMAC(buffer.get_data(),
                                          buffer.get_len(),
                                          (const unsigned char*)pl__key,
                                          pl__key.lengthof());


    return stream_MAC == calculated_MAC;
  }

	TTCN_warning("No message integrity specified.");

	return false;

} //ef__STUN__Check__Message__Integrity

OCT12 ef__STUN__Generate__Transaction__Id() {
  const int transaction_id_length = 12;
  unsigned char id[transaction_id_length];

  if(RAND_pseudo_bytes(id, transaction_id_length) == -1){
    TTCN_error("Not supported by the current RAND method.");
  }

  return OCTETSTRING(transaction_id_length, id);
} //ef__STUN__Generate_Trasaction_Id


INTEGER ef__STUN__get__Message__Length(const OCTETSTRING& pl__stream) {
  int pdu_length = -1;

  if (pl__stream.lengthof() >= 4){
    const unsigned char* stream_pointer = (const unsigned char*)pl__stream;
    pdu_length= (stream_pointer[3] << 8) + stream_pointer[2];
  }

  return pdu_length;
} //ef__STUN__get_Message_Length

void ef__STUN__enc(const PDU__STUN& pl__pdu,
                   const OCT16& pl__key,
                   OCTETSTRING& pl__result,
                   const BOOLEAN& pl__calculate__HMAC__CRC) {
  TTCN_Buffer buf;
  TTCN_EncDec::error_type_t err;
  buf.clear();
  TTCN_EncDec::clear_error();
  TTCN_EncDec::set_error_behavior(TTCN_EncDec::ET_ALL, TTCN_EncDec::EB_WARNING);
  
  pl__pdu.encode(PDU__STUN_descr_, buf, TTCN_EncDec::CT_RAW);

  err = TTCN_EncDec::get_last_error_type();

  if(err != TTCN_EncDec::ET_NONE) {
    TTCN_warning("Encoding error: %s\n", TTCN_EncDec::get_error_str());
  }

  if(pl__calculate__HMAC__CRC) {
    const STUN__Attributes attributes = pl__pdu.STUN__attributes();
    int buffer_position = 0;
    int offset = 0;
    bool insert_hmac = false;
    bool insert_fingerprint = false;
    bool message_length_changed = false;

    if(attributes.lengthof() == 0){
      pl__result = OCTETSTRING(buf.get_len(), buf.get_data());
      return;
    }

    switch(attributes[attributes.lengthof()-1].attribute__type()){
    case STUN__Attribute__Type::MESSAGE__INTEGRITY:
      insert_hmac = true;
      buffer_position = buf.get_len()-MESSAGE_INTEGRITY_LENGTH;
    break;
    case STUN__Attribute__Type::FINGERPRINT:
      insert_fingerprint = true;
      offset = buf.get_len() - FINGERPRINT_LENGTH;

      if(attributes.lengthof() >= 2 && attributes[attributes.lengthof()-2].attribute__type() == STUN__Attribute__Type::MESSAGE__INTEGRITY){
        insert_hmac = true;
        change_Message_Length(buf, (offset-STUN_MESSAGE_HEADER_LENGTH));
        offset -= MESSAGE_INTEGRITY_LENGTH;
        message_length_changed = true;
      }

      buffer_position = offset;
    break;
    default:
      TTCN_warning("When present, the FINGERPRINT attribute MUST be the last attribute in the message, and thus will appear after MESSAGE-INTEGRITY.");
    };

    buf.set_pos(buffer_position);
    buf.cut_end();

    if(insert_hmac){
      insert_HMAC(buf, (const unsigned char*)pl__key, pl__key.lengthof());
    }
    if(insert_fingerprint){
      if(message_length_changed) {
        change_Message_Length(buf, buf.get_len()-STUN_MESSAGE_HEADER_LENGTH+FINGERPRINT_LENGTH);
      }
      insert_Fingerprint(buf);
    }

  }//if calculate hmac and crc

  pl__result = OCTETSTRING(buf.get_len(), buf.get_data());

}//ef__STUN__enc

}//namespace STUN__Types

TTCN_Module STUN_EncDec("STUN_EncDec", __DATE__, __TIME__);