Merge pull request #58 from Zephyruso/ip-device-map

4 files changed, 105 insertions, 6 deletions

diff --git a/src/components/Connections.tsx b/src/components/Connections.tsx
index 7509325..77955c2 100644
--- a/src/components/Connections.tsx
+++ b/src/components/Connections.tsx
@@ -6,6 +6,7 @@ import { useTranslation } from 'react-i18next';
 import { Tab, TabList, TabPanel, Tabs } from 'react-tabs';
 
 import { ConnectionItem } from '~/api/connections';
+import BaseModal from '~/components/shared/BaseModal';
 import { State } from '~/store/types';
 
 import * as connAPI from '../api/connections';
@@ -22,6 +23,10 @@ import SvgYacd from './SvgYacd';
 
 const { useEffect, useState, useRef, useCallback } = React;
 
+const sourceMapInit = localStorage.getItem('sourceMap')
+  ? JSON.parse(localStorage.getItem('sourceMap'))
+  : [];
+
 const paddingBottom = 30;
 
 function arrayToIdKv<T extends { id: string }>(items: T[]) {
@@ -89,14 +94,35 @@ function filterConns(conns: FormattedConn[], keyword: string, sourceIp: string)
   return result;
 }
 
-function getConnIpList(conns: FormattedConn[]) {
-  return Array.from(new Set(conns.map((x) => x.sourceIP))).sort();
+function getConnIpList(conns: FormattedConn[], sourceMap: { reg: string; name: string }[]) {
+  return Array.from(new Set(conns.map((x) => x.sourceIP)))
+    .sort()
+    .map((value) => {
+      return getNameFromSource(value, sourceMap);
+    });
+}
+
+function getNameFromSource(source: string, sourceMap: { reg: string; name: string }[]): string {
+  let sourceName = source;
+
+  sourceMap.forEach(({ reg, name }) => {
+    if (!reg) return;
+
+    const regExp = new RegExp(reg, 'g');
+
+    if (regExp.test(source) && name) {
+      sourceName = `${name}(${source})`;
+    }
+  });
+
+  return sourceName;
 }
 
 function formatConnectionDataItem(
   i: ConnectionItem,
   prevKv: Record<string, { upload: number; download: number }>,
-  now: number
+  now: number,
+  sourceMap: { reg: string; name: string }[]
 ): FormattedConn {
   const { id, metadata, upload, download, start, chains, rule, rulePayload } = i;
   const {
@@ -115,6 +141,8 @@ function formatConnectionDataItem(
   let host2 = host;
   if (host2 === '') host2 = destinationIP;
   const prev = prevKv[id];
+  const source = `${sourceIP}:${sourcePort}`;
+
   const ret = {
     id,
     upload,
@@ -126,7 +154,7 @@ function formatConnectionDataItem(
     host: `${host2}:${destinationPort}`,
     sniffHost: sniffHost ? sniffHost : '-',
     type: `${type}(${network})`,
-    source: `${sourceIP}:${sourcePort}`,
+    source: getNameFromSource(source, sourceMap),
     downloadSpeedCurr: download - (prev ? prev.download : 0),
     uploadSpeedCurr: upload - (prev ? prev.upload : 0),
     process: process ? process : '-',
@@ -168,6 +196,8 @@ function ConnQty({ qty }) {
 }
 
 function Conn({ apiConfig }) {
+  const [sourceMapModal, setSourceMapModal] = useState(false);
+  const [sourceMap, setSourceMap] = useState(sourceMapInit);
   const [refContainer, containerHeight] = useRemainingViewPortHeight();
 
   const [conns, setConns] = useState([]);
@@ -179,7 +209,7 @@ function Conn({ apiConfig }) {
   const filteredConns = filterConns(conns, filterKeyword, filterSourceIpStr);
   const filteredClosedConns = filterConns(closedConns, filterKeyword, filterSourceIpStr);
 
-  const connIpSet = getConnIpList(conns);
+  const connIpSet = getConnIpList(conns, sourceMap);
   // const ClosedConnIpSet = getConnIpList(closedConns);
 
   const [isCloseAllModalOpen, setIsCloseAllModalOpen] = useState(false);
@@ -199,7 +229,7 @@ function Conn({ apiConfig }) {
       const prevConnsKv = arrayToIdKv(prevConnsRef.current);
       const now = Date.now();
       const x = connections.map((c: ConnectionItem) =>
-        formatConnectionDataItem(c, prevConnsKv, now)
+        formatConnectionDataItem(c, prevConnsKv, now, sourceMap)
       );
       const closed = [];
       for (const c of prevConnsRef.current) {
@@ -232,9 +262,65 @@ function Conn({ apiConfig }) {
   }, [apiConfig, read, reConnectCount, setReConnectCount]);
 
   const { t } = useTranslation();
+  const openModalSource = () => {
+    if (sourceMap.length === 0) {
+      sourceMap.push({
+        reg: '',
+        name: '',
+      });
+    }
+    setSourceMapModal(true);
+  };
+  const closeModalSource = () => {
+    setSourceMap(sourceMap.filter((i) => i.reg || i.name));
+    localStorage.setItem('sourceMap', JSON.stringify(sourceMap));
+    setSourceMapModal(false);
+  };
+  const setSource = (key, index, val) => {
+    sourceMap[index][key] = val;
+    setSourceMap(Array.from(sourceMap));
+  };
 
   return (
     <div>
+      <BaseModal isOpen={sourceMapModal} onRequestClose={closeModalSource}>
+        <table>
+          <thead>
+            <tr>
+              <th>{t('c_source')}</th>
+              <th>{t('device_name')}</th>
+            </tr>
+          </thead>
+          <tbody>
+            {sourceMap.map((source, index) => (
+              <tr key={`${index}`}>
+                <td>
+                  <input
+                    type="text"
+                    name="reg"
+                    autoComplete="off"
+                    value={source.reg}
+                    onChange={(e) => setSource('reg', index, e.target.value)}
+                  />
+                </td>
+                <td>
+                  <input
+                    type="text"
+                    name="name"
+                    autoComplete="off"
+                    value={source.name}
+                    onChange={(e) => setSource('name', index, e.target.value)}
+                  />
+                </td>
+                <td>
+                  <Button onClick={() => sourceMap.splice(index, 1)}>{t('delete')}</Button>
+                </td>
+              </tr>
+            ))}
+          </tbody>
+        </table>
+        <Button onClick={() => sourceMap.push({ reg: '', name: '' })}>{t('add_tag')}</Button>
+      </BaseModal>
       <div className={s.header}>
         <ContentHeader title={t('Connections')} />
         <div className={s.inputWrapper}>
@@ -274,6 +360,7 @@ function Conn({ apiConfig }) {
           </TabList>
 
           <div>
+            <Button onClick={openModalSource}>{t('client_tag')}</Button>
             <Button onClick={() => setFilterSourceIpStr('')} kind="minimal">
               {t('All')}
             </Button>
diff --git a/src/i18n/en.ts b/src/i18n/en.ts
index f9e9785..6ca6425 100644
--- a/src/i18n/en.ts
+++ b/src/i18n/en.ts
@@ -68,4 +68,8 @@ export const data = {
   close_all_confirm_no: 'No',
   manage_column: 'Custom columns',
   reset_column: 'Reset columns',
+  device_name: 'Device Tag',
+  delete: 'Delete',
+  add_tag: 'Add tag',
+  client_tag: 'Client tags',
 };
diff --git a/src/i18n/zh-cn.ts b/src/i18n/zh-cn.ts
index 0b47139..bac5ccf 100644
--- a/src/i18n/zh-cn.ts
+++ b/src/i18n/zh-cn.ts
@@ -69,4 +69,8 @@ export const data = {
   close_all_confirm_no: '取消',
   manage_column: '管理列',
   reset_column: '重置列',
+  device_name: '设备名',
+  delete: '删除',
+  add_tag: '添加标签',
+  client_tag: '客户端标签',
 };
diff --git a/src/i18n/zh-tw.ts b/src/i18n/zh-tw.ts
index 4467552..18049d4 100644
--- a/src/i18n/zh-tw.ts
+++ b/src/i18n/zh-tw.ts
@@ -68,4 +68,8 @@ export const data = {
   close_all_confirm_no: '取消',
   manage_column: '管理列',
   reset_column: '重置列',
+  device_name: '设备名',
+  delete: '删除',
+  add_tag: '添加标签',
+  client_tag: '客户端标签',
 };