/*

 * Copyright 2018 Richard Hughes <richard@hughsie.com>

 *

 * SPDX-License-Identifier: LGPL-2.1-or-later

 */

#include "config.h"

#include <string.h>

#include "fu-wacom-usb-firmware.h"

#include "fu-wacom-usb-struct.h"

struct _FuWacomUsbFirmware {

  FuFirmware parent_instance;

};

G_DEFINE_TYPE(FuWacomUsbFirmware, fu_wacom_usb_firmware, FU_TYPE_FIRMWARE)

#define FU_WACOM_USB_FIRMWARE_TOKENS_MAX   100000 /* lines */

#define FU_WACOM_USB_FIRMWARE_SECTIONS_MAX 10

typedef struct {

  guint32 addr;

  guint32 sz;

  guint32 prog_start_addr;

} FuFirmwareWacHeaderRecord;

typedef struct {

  FuFirmware *firmware;

  FuFirmwareParseFlags flags;

  GPtrArray *header_infos;

  GString *image_buffer;

  guint8 images_cnt;

} FuWacomUsbFirmwareTokenHelper;

static gboolean

fu_wacom_usb_firmware_tokenize_cb(GString *token,

          guint token_idx,

          gpointer user_data,

          GError **error)

{

  FuWacomUsbFirmwareTokenHelper *helper = (FuWacomUsbFirmwareTokenHelper *)user_data;

  g_autofree gchar *cmd = NULL;

  /* sanity check */

  if (token_idx > FU_WACOM_USB_FIRMWARE_TOKENS_MAX) {

    g_set_error_literal(error,

            FWUPD_ERROR,

            FWUPD_ERROR_INVALID_DATA,

            "file has too many lines");

    return FALSE;

  }

  /* remove WIN32 line endings */

  g_strdelimit(token->str, "\r\x1a", '\0');

  token->len = strlen(token->str);

  /* ignore blank lines */

  cmd = g_strndup(token->str, 2);

  if (g_strcmp0(cmd, "") == 0)

    return TRUE;

  /* custom metadata */

  if (g_strcmp0(cmd, "WA") == 0) {

    /* header info record */

    if (token->len > 3 && memcmp(token->str + 2, "COM", 3) == 0) {

      guint8 header_image_cnt = 0;

      if (token->len != 40) {

        g_set_error(error,

              FWUPD_ERROR,

              FWUPD_ERROR_INTERNAL,

              "invalid header, got %" G_GSIZE_FORMAT " bytes",

              token->len);

        return FALSE;

      }

      /* sanity check */

      if (helper->header_infos->len > FU_WACOM_USB_FIRMWARE_SECTIONS_MAX) {

        g_set_error(error,

              FWUPD_ERROR,

              FWUPD_ERROR_INTERNAL,

              "too many metadata sections: %u",

              helper->header_infos->len);

        return FALSE;

      }

      if (!fu_firmware_strparse_uint4_safe(token->str,

                   token->len,

                   5,

                   &header_image_cnt,

                   error))

        return FALSE;

      for (guint j = 0; j < header_image_cnt; j++) {

        g_autofree FuFirmwareWacHeaderRecord *hdr = NULL;

        hdr = g_new0(FuFirmwareWacHeaderRecord, 1);

        if (!fu_firmware_strparse_uint32_safe(token->str,

                      token->len,

                      (j * 16) + 6,

                      &hdr->addr,

                      error))

          return FALSE;

        if (!fu_firmware_strparse_uint32_safe(token->str,

                      token->len,

                      (j * 16) + 14,

                      &hdr->sz,

                      error))

          return FALSE;

        g_debug("header_fw%u_addr: 0x%x", j, hdr->addr);

        g_debug("header_fw%u_sz:   0x%x", j, hdr->sz);

        g_ptr_array_add(helper->header_infos, g_steal_pointer(&hdr));

      }

      return TRUE;

    }

    /* firmware headline record */

    if (token->len == 13) {

      FuFirmwareWacHeaderRecord *hdr;

      guint8 idx = 0;

      if (!fu_firmware_strparse_uint4_safe(token->str,

                   token->len,

                   2,

                   &idx,

                   error))

        return FALSE;

      if (idx == 0) {

        g_set_error(error,

              FWUPD_ERROR,

              FWUPD_ERROR_INTERNAL,

              "headline %u invalid",

              idx);

        return FALSE;

      }

      if (idx > helper->header_infos->len) {

        g_set_error(error,

              FWUPD_ERROR,

              FWUPD_ERROR_INTERNAL,

              "headline %u exceeds header count %u",

              idx,

              helper->header_infos->len);

        return FALSE;

      }

      if (idx - 1 != helper->images_cnt) {

        g_set_error(error,

              FWUPD_ERROR,

              FWUPD_ERROR_INTERNAL,

              "headline %u is not in sorted order",

              idx);

        return FALSE;

      }

      hdr = g_ptr_array_index(helper->header_infos, idx - 1);

      if (!fu_firmware_strparse_uint32_safe(token->str,

                    token->len,

                    3,

                    &hdr->prog_start_addr,

                    error))

        return FALSE;

      if (hdr->prog_start_addr != hdr->addr) {

        g_set_error(error,

              FWUPD_ERROR,

              FWUPD_ERROR_INTERNAL,

              "programming address 0x%x != "

              "base address 0x%0x for idx %u",

              hdr->prog_start_addr,

              hdr->addr,

              idx);

        return FALSE;

      }

      g_debug("programing-start-address: 0x%x", hdr->prog_start_addr);

      return TRUE;

    }

    g_debug("unknown Wacom-specific metadata");

    return TRUE;

  }

  /* start */

  if (g_strcmp0(cmd, "S0") == 0) {

    if (helper->image_buffer->len > 0) {

      g_set_error_literal(error,

              FWUPD_ERROR,

              FWUPD_ERROR_INTERNAL,

              "duplicate S0 without S7");

      return FALSE;

    }

    g_string_append_printf(helper->image_buffer, "%s\n", token->str);

    return TRUE;

  }

  /* these are things we want to include in the image */

  if (g_strcmp0(cmd, "S1") == 0 || g_strcmp0(cmd, "S2") == 0 || g_strcmp0(cmd, "S3") == 0 ||

      g_strcmp0(cmd, "S5") == 0 || g_strcmp0(cmd, "S7") == 0 || g_strcmp0(cmd, "S8") == 0 ||

      g_strcmp0(cmd, "S9") == 0) {

    if (helper->image_buffer->len == 0) {

      g_set_error(error, FWUPD_ERROR, FWUPD_ERROR_INTERNAL, "%s without S0", cmd);

      return FALSE;

    }

    g_string_append_printf(helper->image_buffer, "%s\n", token->str);

  } else {

    g_set_error(error,

          FWUPD_ERROR,

          FWUPD_ERROR_INTERNAL,

          "invalid SREC command on line %u: %s",

          token_idx + 1,

          cmd);

    return FALSE;

  }

  /* end */

  if (g_strcmp0(cmd, "S7") == 0) {

    g_autoptr(GBytes) blob = NULL;

    g_autoptr(GBytes) fw_srec = NULL;

    g_autoptr(FuFirmware) firmware_srec = fu_srec_firmware_new();

    g_autoptr(FuFirmware) img = fu_firmware_new();

    FuFirmwareWacHeaderRecord *hdr;

    /* get the correct relocated start address */

    if (helper->images_cnt >= helper->header_infos->len) {

      g_set_error(error,

            FWUPD_ERROR,

            FWUPD_ERROR_INTERNAL,

            "%s without header",

            cmd);

      return FALSE;

    }

    hdr = g_ptr_array_index(helper->header_infos, helper->images_cnt);

    if (helper->image_buffer->len == 0) {

      g_set_error(error,

            FWUPD_ERROR,

            FWUPD_ERROR_INTERNAL,

            "%s with missing image buffer",

            cmd);

      return FALSE;

    }

    /* parse SREC file and add as image */

    blob = g_bytes_new(helper->image_buffer->str, helper->image_buffer->len);

    fu_srec_firmware_set_addr_min(FU_SREC_FIRMWARE(firmware_srec), hdr->addr);

    if (!fu_firmware_parse_bytes(firmware_srec,

               blob,

               0x0,

               helper->flags | FU_FIRMWARE_PARSE_FLAG_NO_SEARCH,

               error))

      return FALSE;

    fw_srec = fu_firmware_get_bytes(firmware_srec, error);

    if (fw_srec == NULL)

      return FALSE;

    fu_firmware_set_bytes(img, fw_srec);

    fu_firmware_set_addr(img, fu_firmware_get_addr(firmware_srec));

    fu_firmware_set_idx(img, helper->images_cnt);

    if (!fu_firmware_add_image(helper->firmware, img, error))

      return FALSE;

    helper->images_cnt++;

    /* clear the image buffer */

    g_string_set_size(helper->image_buffer, 0);

  }

  /* success */

  return TRUE;

}

static gboolean

fu_wacom_usb_firmware_validate(FuFirmware *firmware,

             GInputStream *stream,

             gsize offset,

             GError **error)

{

  return fu_struct_wacom_usb_firmware_hdr_validate_stream(stream, offset, error);

}

static gboolean

fu_wacom_usb_firmware_parse(FuFirmware *firmware,

          GInputStream *stream,

          FuFirmwareParseFlags flags,

          GError **error)

{

  g_autoptr(GPtrArray) header_infos = g_ptr_array_new_with_free_func(g_free);

  g_autoptr(GString) image_buffer = g_string_new(NULL);

  FuWacomUsbFirmwareTokenHelper helper = {.firmware = firmware,

            .flags = flags,

            .header_infos = header_infos,

            .image_buffer = image_buffer,

            .images_cnt = 0};

  /* tokenize */

  if (!fu_strsplit_stream(stream,

        0x0,

        "\n",

        fu_wacom_usb_firmware_tokenize_cb,

        &helper,

        error))

    return FALSE;

  /* verify data is complete */

  if (helper.image_buffer->len > 0) {

    g_set_error_literal(error,

            FWUPD_ERROR,

            FWUPD_ERROR_INTERNAL,

            "truncated data: no S7");

    return FALSE;

  }

  /* ensure this matched the header */

  if (helper.header_infos->len != helper.images_cnt) {

    g_set_error(error,

          FWUPD_ERROR,

          FWUPD_ERROR_INTERNAL,

          "not enough images %u for header count %u",

          helper.images_cnt,

          header_infos->len);

    return FALSE;

  }

  /* success */

  return TRUE;

}

static guint8

fu_wacom_usb_firmware_calc_checksum(GByteArray *buf)

{

  return fu_sum8(buf->data, buf->len) ^ 0xFF;

}

static GByteArray *

fu_wacom_usb_firmware_write(FuFirmware *firmware, GError **error)

{

  g_autoptr(GPtrArray) images = fu_firmware_get_images(firmware);

  g_autoptr(GString) str = g_string_new(NULL);

  g_autoptr(GByteArray) buf = g_byte_array_new();

  g_autoptr(GByteArray) buf_hdr = g_byte_array_new();

  /* fw header */

  if (images->len == 0) {

    g_set_error_literal(error,

            FWUPD_ERROR,

            FWUPD_ERROR_NOT_SUPPORTED,

            "no firmware images found");

    return NULL;

  }

  for (guint i = 0; i < images->len; i++) {

    FuFirmware *img = g_ptr_array_index(images, i);

    fu_byte_array_append_uint32(buf_hdr, fu_firmware_get_addr(img), G_BIG_ENDIAN);

    fu_byte_array_append_uint32(buf_hdr, fu_firmware_get_size(img), G_BIG_ENDIAN);

  }

  g_string_append_printf(str, "WACOM%u", images->len);

  for (guint i = 0; i < buf_hdr->len; i++)

    g_string_append_printf(str, "%02X", buf_hdr->data[i]);

  g_string_append_printf(str, "%02X\n", fu_wacom_usb_firmware_calc_checksum(buf_hdr));

  /* payload */

  for (guint i = 0; i < images->len; i++) {

    FuFirmware *img = g_ptr_array_index(images, i);

    g_autoptr(GBytes) img_blob = NULL;

    g_autoptr(GByteArray) buf_img = g_byte_array_new();

    /* img header */

    g_string_append_printf(str, "WA%u", (guint)fu_firmware_get_idx(img) + 1);

    fu_byte_array_append_uint32(buf_img, fu_firmware_get_addr(img), G_BIG_ENDIAN);

    for (guint j = 0; j < buf_img->len; j++)

      g_string_append_printf(str, "%02X", buf_img->data[j]);

    g_string_append_printf(str, "%02X\n", fu_wacom_usb_firmware_calc_checksum(buf_img));

    /* srec */

    img_blob = fu_firmware_write(img, error);

    if (img_blob == NULL)

      return NULL;

    g_string_append_len(str,

            (const gchar *)g_bytes_get_data(img_blob, NULL),

            g_bytes_get_size(img_blob));

  }

  /* success */

  g_byte_array_append(buf, (const guint8 *)str->str, str->len);

  return g_steal_pointer(&buf);

}

static void

fu_wacom_usb_firmware_init(FuWacomUsbFirmware *self)

{

  fu_firmware_set_images_max(FU_FIRMWARE(self), 1024);

  fu_firmware_add_image_gtype(FU_FIRMWARE(self), FU_TYPE_FIRMWARE);

  fu_firmware_add_image_gtype(FU_FIRMWARE(self), FU_TYPE_SREC_FIRMWARE);

}

static void

fu_wacom_usb_firmware_class_init(FuWacomUsbFirmwareClass *klass)

{

  FuFirmwareClass *firmware_class = FU_FIRMWARE_CLASS(klass);

  firmware_class->validate = fu_wacom_usb_firmware_validate;

  firmware_class->parse = fu_wacom_usb_firmware_parse;

  firmware_class->write = fu_wacom_usb_firmware_write;

}

FuFirmware *

fu_wacom_usb_firmware_new(void)

{

  return FU_FIRMWARE(g_object_new(FU_TYPE_WACOM_USB_FIRMWARE, NULL));

}