link-stack/apps/leafcutter/app/_lib/opensearch.ts

/* eslint-disable no-underscore-dangle */
import { Client } from "@opensearch-project/opensearch";
import { v4 as uuid } from "uuid";

/* Common */

const globalIndex = ".kibana_1";
const dataIndexName = "sample_tagged_tickets";
const userMetadataIndexName = "user_metadata";

const baseURL = `https://${process.env.OPENSEARCH_USERNAME}:${process.env.OPENSEARCH_PASSWORD}@${process.env.OPENSEARCH_URL}`;

const createClient = () => new Client({
  node: baseURL,
  auth: {
    username: process.env.OPENSEARCH_USERNAME!,
    password: process.env.OPENSEARCH_PASSWORD!,
  },
  ssl: {
    rejectUnauthorized: false,
  },
});

const createUserClient = (username: string, password: string) => new Client({
  node: baseURL,
  auth: {
    username,
    password,
  },
  ssl: {
    rejectUnauthorized: false,
  },
});

export const checkAuth = async (username: string, password: string) => {
  const client = createUserClient(username, password);
  const res = await client.ping();

  return res.statusCode === 200;
};

const getDocumentID = (doc: any) => doc._id.split(":")[1];

const getEmbedURL = (tenant: string, visualizationID: string) =>
  `/app/visualize?security_tenant=${tenant}#/edit/${visualizationID}?embed=true`;

export const getVisualization = async (id: string) => {
  const client = createClient();
  const res = await client.get({
    id: `visualization:${id}`,
    index: globalIndex,
  });

  return res.body._source;
};

const generateKuery = (searchQuery: any) => {
  const searchTemplate = {
    query: {
      query: "",
      language: "kuery",
    },
    filter: [],
    indexRefName: "kibanaSavedObjectMeta.searchSourceJSON.index",
  };
  const incidentTypeClause = searchQuery.incidentType.values
    .map((value: string) => `incident:${value} `)
    .join(" or ");
  const allTechnologies = [
    ...searchQuery.platform.values,
    ...searchQuery.device.values,
    ...searchQuery.service.values,
    ...searchQuery.maker.values,
  ];
  const technologyClause = allTechnologies
    .map((value: string) => `technology:${value} `)
    .join(" or ");
  const targetedGroupClause = searchQuery.targetedGroup.values
    .map((value: string) => `targeted_group:${value} `)
    .join(" or ");
  const countryClause = searchQuery.country.values
    .map((value: string) => `country:${value} `)
    .join(" or ");
  const subregionClause = searchQuery.subregion.values
    .map((value: string) => `region:${value} `)
    .join(" or ");
  const continentClause = searchQuery.continent.values
    .map((value: string) => `continent:${value} `)
    .join(" or ");
  const kueryString = [
    incidentTypeClause,
    technologyClause,
    targetedGroupClause,
    countryClause,
    subregionClause,
    continentClause,
  ]
    .filter((clause) => clause !== "")
    .join(" and ");
  searchTemplate.query.query = kueryString;

  return JSON.stringify(searchTemplate);
};

export const getUserMetadata = async (username: string) => {
  const client = createClient();
  let res: any;

  try {
    res = await client.get({
      id: username,
      index: userMetadataIndexName,
    });
  } catch (e) {
    await client.create({
      id: username,
      index: userMetadataIndexName,
      body: { username, savedSearches: [] }
    });

    res = await client.get({
      id: username,
      index: userMetadataIndexName,
    });
  }

  return res?.body._source;
};

export const saveUserMetadata = async (username: string, metadata: any) => {
  const client = createClient();
  await client.update({
    id: username,
    index: userMetadataIndexName,
    body: { doc: { username, ...metadata } }
  });
};

/* User */

const getCurrentUserIndex = async (email: string) => {
  const userIndexName = email.replace(/[\W\d_]/g, "").toLowerCase();
  const client = createClient();
  const aliasesResponse = await client.indices.getAlias({
    name: `.kibana_*_${userIndexName}`,
  });

  // prefer alias if it exists
  if (Object.keys(aliasesResponse.body).length > 0) {
    return Object.keys(aliasesResponse.body)[0];
  }

  const indicesResponse = await client.indices.get({
    index: `.kibana_*_${userIndexName}_1`,
  });
  const currentUserIndex = Object.keys(indicesResponse.body)[0];

  return currentUserIndex;
};

const getIndexPattern: any = async (index: string) => {
  const client = createClient();
  const query = {
    query: {
      bool: {
        must: [
          { match: { type: "index-pattern" } },
          {
            match: {
              "index-pattern.title": dataIndexName,
            },
          },
        ],
      },
    },
  };
  const res = await client.search({
    index,
    size: 1,
    body: query,
    sort: ["updated_at:desc"],
  });

  if (res.body.hits.total.value === 0) {
    // eslint-disable-next-line no-use-before-define
    return createCurrentUserIndexPattern(index);
  }

  const {
    hits: {
      hits: [indexPattern],
    },
  } = res.body;

  return indexPattern;
};

const createCurrentUserIndexPattern = async (index: string) => {
  const { _source: globalIndexPattern } = await getIndexPattern(globalIndex);
  globalIndexPattern.updated_at = new Date().toISOString();

  const id = uuid();
  const fullID = `index-pattern:${id}`;
  const client = createClient();
  const res = await client.create({
    id: fullID,
    index,
    refresh: true,
    body: globalIndexPattern,
  });

  return res.body;
};

const getIndexPatternID = async (index: string) => {
  const indexPattern = await getIndexPattern(index);
  return getDocumentID(indexPattern);
};

interface createUserVisualizationProps {
  email: string;
  query: any;
  visualizationID: string;
  title: string;
  description: string;
}

export const createUserVisualization = async (
  props: createUserVisualizationProps
) => {
  const { email, query, visualizationID, title, description } = props;
  const userIndex = await getCurrentUserIndex(email);
  const indexPatternID = await getIndexPatternID(userIndex);
  const id = uuid();
  const fullID = `visualization:${id}`;

  const template: any = await getVisualization(visualizationID);
  template.visualization.title = title;
  template.visualization.description = description;
  template.visualization.kibanaSavedObjectMeta.searchSourceJSON =
    generateKuery(query);
  template.references = [
    {
      name: "kibanaSavedObjectMeta.searchSourceJSON.index",
      type: "index-pattern",
      id: indexPatternID,
    },
  ];
  template.updated_at = new Date().toISOString();

  const client = createClient();
  const res = await client.create({
    id: fullID,
    index: userIndex,
    refresh: true,
    body: template,
  });

  return getDocumentID(res.body);
};

export const getUserVisualization = async (email: string, id: string) => {
  const userIndex = await getCurrentUserIndex(email);
  const client = createClient();
  const res = await client.get({
    id: `visualization:${id}`,
    index: userIndex,
  });

  return res.body;
};

interface updateVisualizationProps {
  email: string;
  id: string;
  query: any;
  title: string;
  description: string;
}

export const updateUserVisualization = async (
  props: updateVisualizationProps
) => {
  const { email, id, query, title, description } = props;
  const userIndex = await getCurrentUserIndex(email);
  const result: any = await getUserVisualization(email, id);
  const body = {
    doc: result._source,
  };
  body.doc.visualization.title = title;
  body.doc.visualization.description = description;
  body.doc.visualization.kibanaSavedObjectMeta.searchSourceJSON =
    generateKuery(query);

  const client = createClient();
  try {
    await client.update({
      id: `visualization:${id}`,
      index: userIndex,
      body,
    });
  } catch (e) {
    // eslint-disable-next-line no-console
    console.log({ e });
  }

  return id;
};

export const deleteUserVisualization = async (email: string, id: string) => {
  const userIndex = await getCurrentUserIndex(email);
  const client = createClient();
  client.delete({
    id: `visualization:${id}`,
    index: userIndex,
  });
};

export const getUserVisualizations = async (email: string, limit: number) => {
  const userIndex = await getCurrentUserIndex(email);
  const client = createClient();
  const query = {
    query: {
      match: { type: "visualization" },
    },
  };
  const res = await client.search({
    index: userIndex,
    size: limit,
    body: query,
    sort: ["updated_at:desc"],
  });
  const {
    hits: { hits },
  } = res.body;
  const results = hits.map((hit: any) => ({
    id: getDocumentID(hit),
    title: hit._source.visualization.title,
    description: hit._source.visualization.description ?? "",
    url: getEmbedURL("private", getDocumentID(hit)),
  }));

  return results;
};

/* Global */

export const performQuery = async (searchQuery: any, limit: number) => {
  const client = createClient();
  const body = {
    query: {
      bool: {
        must: [],
      },
    },
  };

  if (searchQuery.relativeDate.values.length > 0) {
    searchQuery.relativeDate.values.forEach((value: string) => {
      // @ts-expect-error
      body.query.bool.must.push({
        range: {
          date: {
            gte: `now-${value}d`,
          },
        },
      });
    });
  }

  if (searchQuery.startDate.values.length > 0) {
    searchQuery.startDate.values.forEach((value: string) => {
      // @ts-expect-error
      body.query.bool.must.push({
        range: {
          date: {
            gte: value,
          },
        },
      });
    });
  }

  if (searchQuery.endDate.values.length > 0) {
    searchQuery.endDate.values.forEach((value: string) => {
      // @ts-expect-error
      body.query.bool.must.push({
        range: {
          date: {
            lte: value,
          },
        },
      });
    });
  }

  if (searchQuery.incidentType.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "incident.keyword": searchQuery.incidentType.values },
    });
  }

  if (searchQuery.targetedGroup.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "targeted_group.keyword": searchQuery.targetedGroup.values },
    });
  }

  if (searchQuery.platform.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "technology.keyword": searchQuery.platform.values },
    });
  }

  if (searchQuery.device.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "technology.keyword": searchQuery.device.values },
    });
  }

  if (searchQuery.service.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "technology.keyword": searchQuery.service.values },
    });
  }

  if (searchQuery.maker.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "technology.keyword": searchQuery.maker.values },
    });
  }

  if (searchQuery.subregion.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "region.keyword": searchQuery.subregion.values },
    });
  }

  if (searchQuery.country.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "country.keyword": searchQuery.country.values },
    });
  }

  if (searchQuery.continent.values.length > 0) {
    // @ts-expect-error
    body.query.bool.must.push({
      terms: { "continent.keyword": searchQuery.continent.values },
    });
  }

  const dataResponse = await client.search({
    index: dataIndexName,
    size: limit,
    body,
    sort: ["date:desc"],
  });
  const {
    hits: { hits },
  } = dataResponse.body;
  const results = hits.map((hit: any) => ({
    ...hit._source,
    id: hit._id,
    incident: Array.isArray(hit._source.incident) ? hit._source.incident.join(", ") : hit._source.incident,
    technology: Array.isArray(hit._source.technology) ? hit._source.technology.join(", ") : hit._source.technology,
    targeted_group: Array.isArray(hit._source.targeted_group) ? hit._source.targeted_group.join(", ") : hit._source.targeted_group,
    country: Array.isArray(hit._source.country) ? hit._source.country.join(", ") : hit._source.country,
  }));

  return results;
};

const cleanTitle = (title: string) =>
  title?.replace(/^\[[a-zA-Z]+\] /g, "") ?? "";

const getVisualizationType = (hit: any) => {
  const { title, visState } = hit._source.visualization;
  const rawType = JSON.parse(visState).type;
  let type = "unknown";
  if (
    rawType === "horizontal_bar" &&
    title.includes("Horizontal Bar Stacked")
  ) {
    type = "horizontalBarStacked";
  } else if (rawType === "horizontal_bar" && title.includes("Horizontal Bar")) {
    type = "horizontalBar";
  } else if (
    rawType === "histogram" &&
    title.includes("Vertical Bar Stacked")
  ) {
    type = "verticalBarStacked";
  } else if (rawType === "histogram" && title.includes("Vertical Bar")) {
    type = "verticalBar";
  } else if (rawType === "histogram" && title.includes("Line Stacked")) {
    type = "lineStacked";
  } else if (rawType === "histogram" && title.includes("Line")) {
    type = "line";
  } else if (rawType === "pie") {
    type = "pieDonut";
  } else if (rawType === "table") {
    type = "dataTable";
  } else if (rawType === "metric") {
    type = "metric";
  } else if (rawType === "tagcloud") {
    type = "tagCloud";
  }

  return type;
};

export const getTrends = async (limit: number) => {
  const client = createClient();
  const query = {
    query: {
      bool: {
        must: [
          { match: { type: "visualization" } },
          {
            match_bool_prefix: {
              "visualization.title": "[Trend]",
            },
          },
        ],
      },
    },
  };
  const rawResponse = await client.search({
    index: globalIndex,
    size: limit,
    body: query,
    sort: ["updated_at:desc"],
  });
  const response = rawResponse.body;
  const {
    hits: { hits },
  } = response;
  const results = hits.map((hit: any) => ({
    id: getDocumentID(hit),
    title: cleanTitle(hit._source.visualization.title),
    description: hit._source.visualization.description,
    url: getEmbedURL("global", getDocumentID(hit)),
  }));

  return results;
};

export const getTemplates = async (limit: number) => {
  const client = createClient();

  const query = {
    query: {
      bool: {
        must: [
          { match: { type: "visualization" } },
          {
            match_bool_prefix: {
              "visualization.title": "Templated",
            },
          },
        ],
      },
    },
  };


  const rawResponse = await client.search({
    index: globalIndex,
    size: limit,
    body: query,
  });

  const response = rawResponse.body;
  const {
    hits: { hits },
  } = response;
  const results = hits.map((hit: any) => ({
    id: getDocumentID(hit),
    title: cleanTitle(hit._source.visualization.title),
    type: getVisualizationType(hit),
  }));

  results.sort((a: any, b: any) => a.title.localeCompare(b.title));

  return results;
};