mitm-webserver/frontend/src/components/FirewallRuleBuilder.tsx

// src/components/RuleBuilder.tsx
import { CopyOutlined, PlusOutlined, ReloadOutlined } from '@ant-design/icons';
import {
  Button,
  Card,
  Col,
  Divider,
  Form,
  Input,
  InputNumber,
  message,
  Modal,
  Radio,
  Row,
  Select,
  Space,
  Tooltip,
  Typography,
} from 'antd';
import React, { useCallback, useEffect, useMemo, useState } from 'react';
import { execFirewallRaw } from '../api/apiClient';
import type { ChainOut, ExecResult, Expr, RuleOut, TableOut } from '../types/firewall';

const { Option } = Select;
const { TextArea } = Input;
const { Title, Text } = Typography;

type FormValues = Record<string, any>;

interface RuleBuilderProps {
  onCreated?: () => Promise<void> | void;
  tables?: TableOut[] | null;
  rulesLoading?: boolean;
  rulesError?: string | null;
  refreshRules?: () => Promise<void>;
  onRulesChange?: (tables: TableOut[]) => void;
}

/* ----------------------
   Small helper functions
   ---------------------- */

function isPresetProtocolChoice(values: FormValues) {
  return (values.protocolChoice ?? 'preset') === 'preset';
}

/**
 * Build NFT expression tokens (keeps your original logic).
 */
export function buildExprFromValues(values: FormValues): Expr[] {
  const expr: Expr[] = [];

  const protoChoice = values.protocolChoice ?? 'preset';
  if (protoChoice === 'preset') {
    const preset = values.protocolPreset;
    if (preset === 'icmp') {
      expr.push({
        match: {
          left: { payload: { protocol: 'ip', field: 'protocol' } },
          op: '==',
          right: 'icmp',
        },
      });
    } else if (preset === 'icmpv6') {
      expr.push({
        match: {
          left: { payload: { protocol: 'ip6', field: 'nexthdr' } },
          op: '==',
          right: 'icmpv6',
        },
      });
    } else if (preset === 'tcp') {
      expr.push({ tcp: {} });
    } else if (preset === 'udp') {
      expr.push({ udp: {} });
    }
  } else {
    const custom = (values.protocolCustom || '').trim();
    if (custom) {
      if (/^icmpv6$/i.test(custom)) {
        expr.push({
          match: {
            left: { payload: { protocol: 'ip6', field: 'nexthdr' } },
            op: '==',
            right: 'icmpv6',
          },
        });
      } else if (/^icmp$/i.test(custom)) {
        expr.push({
          match: {
            left: { payload: { protocol: 'ip', field: 'protocol' } },
            op: '==',
            right: 'icmp',
          },
        });
      } else if (/tcp/i.test(custom)) {
        expr.push({ tcp: {} });
      } else if (/udp/i.test(custom)) {
        expr.push({ udp: {} });
      } else {
        try {
          const parsed = JSON.parse(custom);
          expr.push(parsed);
        } catch {
          expr.push(custom);
        }
      }
    }
  }

  if (values.saddr) {
    expr.push({
      match: {
        left: { payload: { protocol: 'ip', field: 'saddr' } },
        op: '==',
        right: values.saddr,
      },
    });
  }
  if (values.daddr) {
    expr.push({
      match: {
        left: { payload: { protocol: 'ip', field: 'daddr' } },
        op: '==',
        right: values.daddr,
      },
    });
  }

  const chooseUdp = (preset?: string, custom?: string) => preset === 'udp' || (custom && /(udp)/i.test(custom));

  if (values.sport) {
    const proto = chooseUdp(values.protocolPreset, values.protocolCustom) ? 'udp' : 'tcp';
    expr.push({ [proto]: { sport: Number(values.sport) } } as any);
  }
  if (values.dport) {
    const proto = chooseUdp(values.protocolPreset, values.protocolCustom) ? 'udp' : 'tcp';
    expr.push({ [proto]: { dport: Number(values.dport) } } as any);
  }

  if (values.advanced) {
    const adv = (values.advanced || '').trim();
    try {
      expr.push(JSON.parse(adv));
    } catch {
      expr.push(adv);
    }
  }

  const action = values.action || 'drop';
  if (action === 'drop') expr.push({ drop: null });
  else if (action === 'accept') expr.push({ accept: null });
  else if (action === 'reject') expr.push({ reject: null });
  else if (action === 'nfqueue' || action === 'queue') {
    const qnum = Number(values.nfqueue || 0);
    expr.push({ queue: Number.isFinite(qnum) && qnum > 0 ? qnum : 0 } as any);
  }

  return expr;
}

/**
 * Create a compact textual representation of an expr.
 */
export function textFromExpr(expr: Expr | null | undefined): string {
  if (!expr) return '';
  if (!Array.isArray(expr)) return String(expr);

  const tokens: string[] = [];
  for (const part of expr) {
    if (typeof part === 'string') {
      tokens.push(part);
      continue;
    }
    if (!part || typeof part !== 'object') {
      tokens.push(String(part));
      continue;
    }
    if ('match' in part) {
      const m = (part as any).match;
      const left = m.left;
      const right = m.right;
      if (left && left.payload && right !== undefined) {
        const p = left.payload;
        const prot = p.protocol;
        const field = p.field;
        if (prot && field && typeof right === 'string') {
          tokens.push(`${prot} ${field} ${right}`);
          continue;
        }
      }
      tokens.push('match');
      continue;
    }
    if ('payload' in part) {
      const p = (part as any).payload;
      tokens.push(`payload(${p.protocol}.${p.field})`);
      continue;
    }
    if ('tcp' in part) {
      const v = (part as any).tcp;
      if (v && v.dport) tokens.push(`tcp dport ${v.dport}`);
      else if (v && v.sport) tokens.push(`tcp sport ${v.sport}`);
      else tokens.push('tcp');
      continue;
    }
    if ('udp' in part) {
      const v = (part as any).udp;
      if (v && v.dport) tokens.push(`udp dport ${v.dport}`);
      else if (v && v.sport) tokens.push(`udp sport ${v.sport}`);
      else tokens.push('udp');
      continue;
    }
    if ('drop' in part) {
      tokens.push('drop');
      continue;
    }
    if ('accept' in part) {
      tokens.push('accept');
      continue;
    }
    if ('reject' in part) {
      tokens.push('reject');
      continue;
    }
    if ('counter' in part) {
      tokens.push('counter');
      continue;
    }
    if ('queue' in part) {
      tokens.push(`queue ${(part as any).queue}`);
      continue;
    }
    tokens.push(Object.keys(part).sort().join('+'));
  }
  return tokens.join(' ');
}

/* -------------------------
   Command generation helper
   ------------------------- */

function generateCommandFromValues(values: FormValues): { cmd: string; textual: string; position?: number | null } {
  const expr = buildExprFromValues(values);
  let textual = textFromExpr(expr).trim();

  if (values.advanced && typeof values.advanced === 'string' && values.advanced.trim() !== '') {
    textual = `${textual} ${values.advanced.trim()}`.trim();
  }

  if (values.action === 'nfqueue' || values.action === 'queue') {
    const qnum = values.nfqueue ?? values.queue ?? null;
    const qstr = qnum != null && String(qnum).trim() !== '' ? String(qnum).trim() : null;
    const hasQueue = /\bqueue(?:\s+num)?\b/i.test(textual);
    if (qstr) {
      if (!hasQueue) textual = `${textual} queue num ${qstr}`.trim();
      else textual = textual.replace(/\bqueue(?:\s+num)?\s*\d*/i, `queue num ${qstr}`);
    } else {
      if (!hasQueue) textual = `${textual} queue num 0`.trim();
      else textual = textual.replace(/\bqueue(?:\s+num)?\s*\d*/i, `queue num 0`);
    }
  }

  const tableSelect = values.tableSelect;
  const chain = values.chainSelect || 'input';
  const [family = 'inet', table = 'filter'] = tableSelect ? String(tableSelect).split(':') : ['inet', 'filter'];

  const before = values.insertBeforeHandle;
  const hasBefore = before != null && String(before) !== '';
  const verb = hasBefore ? 'insert' : 'add';
  const positionPart = hasBefore ? ` position ${before}` : '';

  const cmd = `${verb} rule ${family} ${table} ${chain}${positionPart} ${textual}`.replace(/\s+/g, ' ').trim();

  return { cmd, textual, position: hasBefore ? Number(before) : undefined };
}

/* -------------------------
   Component
   ------------------------- */

export const RuleBuilder: React.FC<RuleBuilderProps> = (props: RuleBuilderProps) => {
  const [form] = Form.useForm();
  const [refreshing, setRefreshing] = useState(false);
  const [loading, setLoading] = useState(false);

  const [cmdPreview, setCmdPreview] = useState('');

  const [insertBeforeOptions, setInsertBeforeOptions] = useState<Array<{ value: any; label: string }>>([]);

  const tableOptions = useMemo(
    () => (props.tables || []).map((t) => ({ value: `${t.family}:${t.name}`, label: `${t.family}:${t.name}` })),
    [props.tables],
  );

  const noTables = !(props.tables && props.tables.length > 0);

  const updateInsertOptions = useCallback(() => {
    const ts = form.getFieldValue('tableSelect');
    const cs = form.getFieldValue('chainSelect');
    if (!ts || !cs) {
      setInsertBeforeOptions([]);
      return;
    }
    const [family, table] = String(ts).split(':');
    const tbl = props.tables?.find((t) => t.family === family && t.name === table);
    if (!tbl) {
      setInsertBeforeOptions([]);
      return;
    }
    const ch = (tbl.chains || []).find((c: ChainOut) => c.name === cs);
    if (!ch || !Array.isArray(ch.rules)) {
      setInsertBeforeOptions([]);
      return;
    }

    const opts = ch.rules
      .filter((r: RuleOut) => r && r.handle != null)
      .map((r: RuleOut) => ({
        value: r.handle,
        label: `#${r.handle} — ${
          r.text ?? (typeof r.expr === 'string' ? r.expr : JSON.stringify(r.expr || {}).slice(0, 120))
        }`,
      }));
    setInsertBeforeOptions(opts);
  }, [form, props.tables]);

  const updatePreviewFromValues = useCallback(
    (values?: FormValues) => {
      const vals = values ?? form.getFieldsValue();
      const { cmd } = generateCommandFromValues(vals);
      setCmdPreview(cmd);
    },
    [form],
  );

  useEffect(() => {
    if (tableOptions.length > 0) {
      const first = tableOptions[0].value;
      form.setFieldsValue({
        tableSelect: first,
        protocolChoice: 'preset',
        protocolPreset: 'icmp',
        action: 'drop',
        nfqueue: 1,
      });

      const [f, n] = String(first).split(':');
      const tbl = props.tables?.find((t) => t.family === f && t.name === n);
      if (tbl && tbl.chains && tbl.chains.length > 0) {
        form.setFieldsValue({ chainSelect: tbl.chains[0].name });
      } else {
        form.setFieldsValue({ chainSelect: undefined });
      }

      setTimeout(() => {
        updateInsertOptions();
        updatePreviewFromValues();
      }, 0);
    } else {
      form.setFieldsValue({
        protocolChoice: 'preset',
        protocolPreset: 'icmp',
        action: 'drop',
        nfqueue: 1,
        tableSelect: undefined,
        chainSelect: undefined,
      });
      setInsertBeforeOptions([]);
      setTimeout(() => updatePreviewFromValues(), 0);
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [props.tables, tableOptions.length]);

  useEffect(() => {
    const current = form.getFieldsValue();
    if (current.nfqueue == null) form.setFieldsValue({ nfqueue: 1 });
    updatePreviewFromValues();
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, []);

  const onValuesChange = useCallback(
    (_: any, allValues: FormValues) => {
      // If user switches to NFQUEUE and nfqueue is empty, seed it with 1
      if (allValues.action === 'nfqueue' && (allValues.nfqueue == null || allValues.nfqueue === '')) {
        form.setFieldsValue({ nfqueue: 1 });
        allValues.nfqueue = 1;
      }

      const ts = allValues.tableSelect;
      if (ts) {
        const [f, n] = String(ts).split(':');
        const tbl = props.tables?.find((t) => t.family === f && t.name === n);
        if (tbl) {
          if (tbl.chains && tbl.chains.length > 0) {
            if (!allValues.chainSelect) {
              form.setFieldsValue({ chainSelect: tbl.chains[0].name });
            }
          } else {
            form.setFieldsValue({ chainSelect: undefined });
          }
        }
      }

      updateInsertOptions();
      updatePreviewFromValues(allValues);
    },
    [form, props.tables, updateInsertOptions, updatePreviewFromValues],
  );

  const handleCreate = useCallback(
    async (values: FormValues) => {
      const { cmd } = generateCommandFromValues(values);

      Modal.confirm({
        title: 'Run raw nft command',
        content: (
          <div>
            <Text>
              About to run nft command in <b>{String(values.tableSelect ?? 'inet:filter')}</b> (see preview).
            </Text>
            <Divider />
            <Text strong>Command:</Text>
            <pre style={{ whiteSpace: 'pre-wrap', marginTop: 8 }}>{cmd}</pre>
          </div>
        ),
        okText: 'Run',
        onOk: async () => {
          setLoading(true);
          try {
            const out: ExecResult = await execFirewallRaw(cmd);
            const stderrText = out?.stderr ? String(out.stderr).trim() : '';
            if (stderrText) {
              message.error(`Command produced stderr: ${stderrText}`);
            } else if (out && (out.rc === 0 || out.rc === -1)) {
              message.success('Command executed (rc ' + out.rc + ')');
              if (props.refreshRules) await props.refreshRules();
              if (props.onCreated) await props.onCreated();
            } else {
              const info = out
                ? `rc:${out.rc}` +
                  (out.stdout ? ` stdout:${out.stdout}` : '') +
                  (out.stderr ? ` stderr:${out.stderr}` : '')
                : 'unknown result';
              message.error('Command failed: ' + info);
            }
          } catch (err: any) {
            message.error('Execution failed: ' + (err?.message ?? String(err)));
          } finally {
            setLoading(false);
          }
        },
      });
    },
    [props.refreshRules, props.onCreated],
  );

  const chainOptions = useMemo(() => {
    const ts = form.getFieldValue('tableSelect');
    if (!ts) return [];
    const [f, n] = String(ts).split(':');
    const tbl = props.tables?.find((t) => t.family === f && t.name === n);
    if (!tbl) return [];
    return tbl.chains.map((c) => (
      <Option key={c.name} value={c.name}>
        {c.name}
      </Option>
    ));
  }, [form, props.tables]);

  const handleRefresh = useCallback(async () => {
    setRefreshing(true);
    try {
      if (props.refreshRules) {
        await props.refreshRules();
        message.success('Rules refresh requested');
      } else {
        message.info('No refresh function provided by parent.');
      }
    } catch (err) {
      console.warn('refresh failed', err);
      message.error('Refresh failed');
    } finally {
      updateInsertOptions();
      setRefreshing(false);
    }
  }, [props.refreshRules, updateInsertOptions]);

  return (
    <Card style={{ maxWidth: 980 }}>
      <Row justify="space-between" align="middle">
        <Col>
          <Title level={4} style={{ margin: 0 }}>
            Add Firewall Rule (raw)
          </Title>
          <Text type="secondary">Builds a raw nft command and executes it via the raw endpoint.</Text>
        </Col>

        <Col>
          <Space>
            <Button icon={<ReloadOutlined />} onClick={handleRefresh} loading={refreshing}>
              Refresh tables
            </Button>
          </Space>
        </Col>
      </Row>

      <Divider />

      <Form
        layout="vertical"
        form={form}
        initialValues={{
          protocolChoice: 'preset',
          protocolPreset: 'icmp',
          action: 'drop',
          nfqueue: 1,
          tableSelect: tableOptions.length > 0 ? tableOptions[0].value : undefined,
        }}
        onFinish={handleCreate}
        onValuesChange={onValuesChange}
      >
        {/* Table & Chain selection */}
        <Row gutter={16}>
          <Col xs={24} sm={12}>
            <Form.Item name="tableSelect" label="Table (family:name)" rules={[{ required: true }]}>
              <Select placeholder={noTables ? 'No tables available' : 'Select table'} disabled={noTables}>
                {tableOptions.map((t) => (
                  <Option key={t.value} value={t.value}>
                    {t.label}
                  </Option>
                ))}
              </Select>
            </Form.Item>
          </Col>

          <Col xs={24} sm={12}>
            <Form.Item name="chainSelect" label="Chain (select)" rules={[{ required: true }]}>
              <Select placeholder="Select chain" disabled={noTables || chainOptions.length === 0}>
                {chainOptions}
              </Select>
            </Form.Item>
          </Col>
        </Row>

        {/* Insert-before */}
        <Row gutter={16}>
          <Col xs={24} sm={12}>
            <Form.Item
              name="insertBeforeHandle"
              label="Insert before (optional)"
              help="Pick an existing rule handle to insert the new rule before. If left empty the rule will be appended."
            >
              <Select allowClear placeholder="Append (no insert-before)">
                {insertBeforeOptions.length === 0 ? (
                  <Option value="__none__" disabled>
                    (no rules available)
                  </Option>
                ) : (
                  insertBeforeOptions.map((o) => (
                    <Option key={String(o.value)} value={o.value}>
                      {o.label}
                    </Option>
                  ))
                )}
              </Select>
            </Form.Item>
          </Col>

          <Col xs={24} sm={12} style={{ display: 'flex', alignItems: 'flex-end' }}>
            <Text type="secondary">
              Use when you want the new rule to appear before a known handle. Refresh tables to see latest handles.
            </Text>
          </Col>
        </Row>

        {/* Protocol / addresses */}
        <Row gutter={16}>
          <Col xs={24} sm={8}>
            <Form.Item label="Protocol selection" style={{ marginBottom: 0 }}>
              <Space>
                <Form.Item name="protocolChoice" noStyle>
                  <Radio.Group>
                    <Radio.Button value="preset">Preset</Radio.Button>
                    <Radio.Button value="custom">Custom</Radio.Button>
                  </Radio.Group>
                </Form.Item>
              </Space>
            </Form.Item>

            {isPresetProtocolChoice(form.getFieldsValue()) ? (
              <Form.Item name="protocolPreset" label="Protocol (preset)">
                <Select>
                  <Option value="">(none)</Option>
                  <Option value="icmp">ICMP (IPv4)</Option>
                  <Option value="icmpv6">ICMPv6</Option>
                  <Option value="tcp">TCP</Option>
                  <Option value="udp">UDP</Option>
                </Select>
              </Form.Item>
            ) : (
              <Form.Item name="protocolCustom" label="Protocol (custom)">
                <Input placeholder="e.g. icmp, icmpv6, tcp, udp, meta ..." />
              </Form.Item>
            )}
          </Col>

          <Col xs={24} sm={8}>
            <Form.Item name="saddr" label="Source IP / CIDR">
              <Input placeholder="10.0.0.0/8 or 192.168.1.5" />
            </Form.Item>
          </Col>

          <Col xs={24} sm={8}>
            <Form.Item name="daddr" label="Destination IP / CIDR">
              <Input placeholder="172.16.0.0/12" />
            </Form.Item>
          </Col>
        </Row>

        {/* Ports / action */}
        <Row gutter={16}>
          <Col xs={24} sm={6}>
            <Form.Item name="sport" label="Source Port">
              <InputNumber style={{ width: '100%' }} min={1} max={65535} />
            </Form.Item>
          </Col>

          <Col xs={24} sm={6}>
            <Form.Item name="dport" label="Dest Port">
              <InputNumber style={{ width: '100%' }} min={1} max={65535} />
            </Form.Item>
          </Col>

          <Col xs={24} sm={12}>
            <Form.Item name="action" label="Action" rules={[{ required: true }]} initialValue="drop">
              <Radio.Group>
                <Radio value="drop">Drop</Radio>
                <Radio value="accept">Accept</Radio>
                <Radio value="reject">Reject</Radio>
                <Radio value="nfqueue">NFQUEUE</Radio>
              </Radio.Group>
            </Form.Item>

            {/* NFQUEUE: required when selected, default 1 */}
            {form.getFieldValue('action') === 'nfqueue' && (
              <Form.Item
                name="nfqueue"
                label="Queue number"
                rules={[
                  { required: true, message: 'Queue number is required when NFQUEUE is selected' },
                  {
                    type: 'number',
                    min: 1,
                    message: 'Queue number must be at least 1',
                  },
                ]}
              >
                <InputNumber style={{ width: '100%' }} min={1} max={65535} />
              </Form.Item>
            )}
          </Col>
        </Row>

        {/* Advanced */}
        <Row>
          <Col span={24}>
            <Form.Item name="advanced" label="Advanced (optional)">
              <TextArea
                rows={3}
                placeholder='Append advanced text (will be appended to the command preview). Example: "limit rate 5/second"'
              />
              <Text type="secondary">
                Advanced is appended to the textual expr preview. Prefer structured fields where possible.
              </Text>
            </Form.Item>
          </Col>
        </Row>

        <Divider />

        <Form.Item label="Preview">
          <Space style={{ display: 'flex', justifyContent: 'space-between', alignItems: 'center' }}>
            <div style={{ flex: 1, marginRight: 12 }}>
              <Text strong>nft command preview:</Text>
              <div style={{ marginTop: 6 }}>
                <Text code copyable={{ text: cmdPreview || '' }}>
                  {cmdPreview || '(no textual preview available)'}
                </Text>
              </div>
            </div>

            <div style={{ width: 320 }}>
              <Tooltip title="Copy command to clipboard">
                <Button
                  icon={<CopyOutlined />}
                  size="small"
                  style={{ marginRight: 8 }}
                  onClick={() => {
                    if (!cmdPreview) return;
                    navigator.clipboard?.writeText(cmdPreview);
                    message.success('Command copied');
                  }}
                />
              </Tooltip>

              <Button type="primary" htmlType="submit" loading={loading} icon={<PlusOutlined />} disabled={noTables}>
                Run command
              </Button>
            </div>
          </Space>
        </Form.Item>
      </Form>
    </Card>
  );
};

export default RuleBuilder;