/**
Copyright: 2024 Contrast Security, Inc
  Contact: support@contrastsecurity.com
  License: Commercial

  NOTICE: This Software and the patented inventions embodied within may only be
  used as part of Contrast Security’s commercial offerings. Even though it is
  made available through public repositories, use of this Software is subject to
  the applicable End User Licensing Agreement found at
  https://www.contrastsecurity.com/enduser-terms-0317a or as otherwise agreed
  between Contrast Security and the End User. The Software may not be reverse
  engineered, modified, repackaged, sold, redistributed or otherwise used in a
  way not consistent with the End User License Agreement.
*/
'use strict';

const agentEmitter = require('../../agent-emitter');
const { PATCH_TYPES } = require('../../constants');
const ModuleHook = require('../../hooks/require');
const patcher = require('../../hooks/patcher');
const logger = require('../logger')('contrast:arch-component');

// Architecture component for versions <3.0.0
ModuleHook.resolve(
  {
    name: 'mongodb',
    file: 'lib/mongo_client.js',
    version: '<3.0.0'
  },
  (MongoClient) => {
    patcher.patch(MongoClient, 'connect', {
      name: 'MongoClient.connect.arch_component',
      patchType: PATCH_TYPES.ARCH_COMPONENT,
      alwaysRun: true,
      pre(ctx) {
        // check if typeof callback == 'function'
        // if yes:
        // - MongoClient.connect executes a cb function, which has access to the connection status
        // if not:
        // - MongoClient.connect should return a promise and we can check it's result in another hook, just as before
        const callbackIndex = ctx.args[2] ? 2 : 1;
        if (ctx.args[callbackIndex] instanceof Function || typeof ctx.args[callbackIndex] === 'function') {
          ctx.args[callbackIndex] = patcher.patch(ctx.args[callbackIndex], {
            name: 'MongoClient.connect.callback.arch_component',
            patchType: PATCH_TYPES.ARCH_COMPONENT,
            alwaysRun: true,
            pre(ctx) {
              const [, db] = ctx.args;
              if (db && db.s.topology && db.s.topology.s) {
                try {
                  const server = db.s.topology.s.server.s;
                  if (server.pool && server.pool.state == 'connected') {
                    const connections = server.pool.availableConnections;
                    for (const c of connections) {
                      agentEmitter.emit('architectureComponent', {
                        vendor: 'MongoDB',
                        url: `mongodb://${c.host}:${c.port}`,
                        remoteHost: '',
                        remotePort: c.port
                      });
                    }
                  }
                } catch (err) {
                  logger.warn('unable to report MongoDB architecture component, err: %o', err);
                }
              }
            }
          });
        }
      },
      post(ctx) {
        if (!ctx.result || !ctx.result.then) {
          return;
        }
        // it never gets here if callbacks are used, because the result won't be then-able
        ctx.result.then(db => {
          if (db && db.s && db.s.topology && db.s.topology.s) {
            try {
              const server = db.s.topology.s.server.s;
              if (server.pool && server.pool.state == 'connected') {
                const connections = server.pool.availableConnections;
                for (const c of connections) {
                  agentEmitter.emit('architectureComponent', {
                    vendor: 'MongoDB',
                    url: `mongodb://${c.host}:${c.port}`,
                    remoteHost: '',
                    remotePort: c.port
                  });
                }
              }
            } catch (err) {
              logger.warn(
                'unable to report MongoDB architecture component, err: %o',
                err,
              );
            }
          }
        });
      }
    });
  }
);

// Architecture component for versions >=3.3.0 <4.0.0
ModuleHook.resolve(
  {
    name: 'mongodb',
    file: 'lib/mongo_client.js',
    version: '>=3.0.0 <4.0.0'
  },
  (MongoClient, { version }) => {
    patcher.patch(MongoClient.prototype, 'connect', {
      name: 'MongoClient.connect.arch_component',
      patchType: PATCH_TYPES.ARCH_COMPONENT,
      alwaysRun: true,
      post(ctx) {
        if (!ctx.result || !ctx.result.then) {
          return;
        }

        // We should report only when connection is successful
        ctx.result.then(function(client) {
          try {
            const { servers = [] } = ctx.obj.s && ctx.obj.s.options;
            for (const server of servers) {
              agentEmitter.emit('architectureComponent', {
                vendor: 'MongoDB',
                url: `mongodb://${server.host}:${server.port}`,
                remoteHost: '',
                remotePort: server.port,
              });
            }
          } catch (err) {
            logger.warn(
              'unable to report MongoDB architecture component, err: %o',
              err,
            );
          }
        });
      },
    });
  },
);

// Architecture component for versions >=4.0.0
ModuleHook.resolve(
  {
    name: 'mongodb',
    version: '>=4.0.0'
  },
  (MongoDB, { version }) => {
    patcher.patch(MongoDB.MongoClient.prototype, 'connect', {
      name: 'MongoClient.connect.arch_component',
      patchType: PATCH_TYPES.ARCH_COMPONENT,
      alwaysRun: true,
      post(ctx) {
        if (!ctx.result || !ctx.result.then) {
          return;
        }

        // We should report only when connection is successful
        ctx.result.then(function(client) {
          if (client && client.topology && client.topology.s) {
            try {
              const { servers } = client.topology.s;
              for (const [, server] of servers) {
                if (server.s && server.s.state === 'connected') {
                  const { srvServiceName } = server.s.options;
                  const { address } = server.s.description;
                  agentEmitter.emit('architectureComponent', {
                    vendor: 'MongoDB',
                    url: `${srvServiceName}://${address}`,
                    remoteHost: '',
                    remotePort: address.split(':').pop()
                  });
                }
              }
            } catch (err) {
              logger.warn(
                'unable to report MongoDB architecture component, err: %o',
                err,
              );
            }
          }
        });
      },
    });
  },
);
