package com.sunyard.snmp;

import org.snmp4j.*;
import org.snmp4j.event.ResponseEvent;
import org.snmp4j.event.ResponseListener;
import org.snmp4j.mp.*;
import org.snmp4j.security.*;
import org.snmp4j.smi.*;
import org.snmp4j.transport.DefaultUdpTransportMapping;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Properties;
import java.util.Vector;

public class SnmpTrapSender implements AutoCloseable{
    private Snmp snmp;

    // 配置参数
    private String ip;
    private int port;
    private String enterpriseOid;
    private String community;
    private int version;

    public SnmpTrapSender(Map<String, String> map) throws Exception {
        loadConfig(map);
        initializeSNMP();
    }

    public void loadConfig(Map<String, String> map) {
        ip = map.getOrDefault("snmp.target.ip", null);
        port = Integer.parseInt(map.getOrDefault("snmp.target.port", "162"));
        enterpriseOid = map.getOrDefault("enterprise.oid", null);
        community = map.getOrDefault("snmp.community", "public");
        version = Integer.parseInt(map.getOrDefault("snmp.version", "1"));
    }

    private void initializeSNMP() throws Exception {
        TransportMapping<?> transport = new DefaultUdpTransportMapping();
        snmp = new Snmp(transport);
        snmp.listen();
    }

    public ResponseEvent sendTrap(String message, int severity) throws Exception {
        ResponseEvent event = null;
        switch (version) {
            case SnmpConstants.version1:
                event = sendV1Trap(message, severity);
                break;
            case SnmpConstants.version2c:
                event = sendV2Trap(message, severity);
                break;
            default:
                throw new IllegalArgumentException("不支持的SNMP版本");
        }
        return event;
    }

    private ResponseEvent sendV1Trap(String message, int severity) throws Exception {
        CommunityTarget target = createCommunityTarget(SnmpConstants.version1);
        PDUv1 pdu = new PDUv1();
        configureCommonTrap(pdu, message, severity);
        pdu.setGenericTrap(PDUv1.ENTERPRISE_SPECIFIC);
        return snmp.send(pdu, target);
    }

    private ResponseEvent sendV2Trap(String message, int severity) throws Exception {
        CommunityTarget target = createCommunityTarget(SnmpConstants.version2c);
        PDU pdu = new PDU();
        configureCommonTrap(pdu, message, severity);
        pdu.setType(PDU.TRAP);
        return snmp.send(pdu, target);
    }


    private void configureCommonTrap(PDU pdu, String message, int severity) {
        pdu.add(new VariableBinding(SnmpConstants.sysUpTime,
                new TimeTicks(System.currentTimeMillis()/1000)));
        pdu.add(new VariableBinding(SnmpConstants.snmpTrapOID,
                new OID(enterpriseOid + ".1.1")));
        byte[] msg = new byte[10];
        try {
            msg = message.getBytes("GBK");
        } catch (Exception e) {
            System.out.println("编码错误");
        }
        pdu.add(new VariableBinding(new OID(enterpriseOid + ".2.1"),
                new OctetString(msg)));
        pdu.add(new VariableBinding(new OID(enterpriseOid + ".2.2"),
                new Integer32(severity)));
    }

    private CommunityTarget createCommunityTarget(int version) {
        CommunityTarget target = new CommunityTarget();
        target.setCommunity(new OctetString(community));
        target.setAddress(GenericAddress.parse("udp:"+ip+"/"+port));
        target.setVersion(version);
        target.setTimeout(5000);
        target.setRetries(3);
        return target;
    }

    @Override
    public void close() throws Exception {
        if (snmp != null) {
            snmp.close();
            snmp = null;
        }
    }
}