diff options
| author | 王宽 <[email protected]> | 2024-02-01 10:42:29 +0000 |
|---|---|---|
| committer | 王宽 <[email protected]> | 2024-02-01 10:42:29 +0000 |
| commit | 9a9fb185e9ee02f2512afab438c7e28c58d0a0aa (patch) | |
| tree | 447c0b6315b778a65f24dd7bffabfccb6ab42774 | |
| parent | 7719a56a565d77b7984057e9bf6e2542eadf27f0 (diff) | |
| parent | bd58ec4f16d001d8088a2f114da6078d0f888fe8 (diff) | |
Merge branch 'feature/cn' into 'develop'
[improve][core]新增CN函数
See merge request galaxy/platform/groot-stream!16
58 files changed, 5053 insertions, 0 deletions
diff --git a/config/cn/cn_grootstream_job_local_template.yaml b/config/cn/cn_grootstream_job_local_template.yaml new file mode 100644 index 0000000..5a8fcb0 --- /dev/null +++ b/config/cn/cn_grootstream_job_local_template.yaml @@ -0,0 +1,316 @@ +sources: + kafka_source: + type: kafka + # fields: # [array of object] Field List, if not set, all fields(Map<String, Object>) will be output. + properties: # [object] Source Properties + topic: SESSION-RECORD + kafka.bootstrap.servers: 192.168.44.11:9094,192.168.44.13:9094,192.168.44.14:9094,192.168.44.15:9094,192.168.44.16:9094 + kafka.session.timeout.ms: 60000 + kafka.max.poll.records: 3000 + kafka.max.partition.fetch.bytes: 31457280 + kafka.security.protocol: SASL_PLAINTEXT + kafka.ssl.keystore.location: + kafka.ssl.keystore.password: + kafka.ssl.truststore.location: + kafka.ssl.truststore.password: + kafka.ssl.key.password: + kafka.sasl.mechanism: PLAIN + kafka.sasl.jaas.config: org.apache.kafka.common.security.plain.PlainLoginModule required username="admin" password="galaxy2019"; + kafka.buffer.memory: + kafka.group.id: local-test + kafka.auto.offset.reset: latest + kafka.max.request.size: + kafka.compression.type: none + format: json # [string] Data Format, default is json + +processing_pipelines: + session_record_processor: # [object] Processing Pipeline + type: com.geedgenetworks.core.processor.projection.ProjectionProcessorImpl + remove_fields: + output_fields: + functions: # [array of object] Function List + - function: SNOWFLAKE_ID + lookup_fields: [ '' ] + output_fields: [ log_id ] + filter: + parameters: + data_center_id_num: 1 + + - function: UNIX_TIMESTAMP_CONVERTER + lookup_fields: [ __timestamp ] + output_fields: [ recv_time ] + parameters: + precision: seconds + + - function: EVAL + output_fields: [ ingestion_time ] + parameters: + value_expression: recv_time + + - function: EVAL + output_fields: [ domain ] + parameters: + value_expression: server_fqdn + + - function: EVAL + output_fields: [ domain_sld ] + parameters: + value_expression: server_domain + + - function: CN_L7_PROTOCOL_AND_APP_EXTRACT + parameters: + decoded_path_field_name: decoded_path + app_transition_field_name: app_transition + l7_protocol_field_name: l7_protocol + app_field_name: app + l7_protocol: DHCP,DNS,FTP,GRE,GTP,HTTP,HTTPS,ICMP,IMAP,IMAPS,IPSEC,ISAKMP,XMPP,L2TP,LDAP,MMS,NETBIOS,NETFLOW,NTP,POP3,POP3S,RDP,PPTP,RADIUS,RTCP,RTP,RTSP,SIP,SMB,SMTP,SMTPS,SNMP,SSDP,SSH,SSL,STUN,TELNET,TFTP,OPENVPN,RTMP,TEREDO,FTPS,DTLS,SPDY,BJNP,QUIC,MDNS,Unknown TCP,Unknown UDP,Unknown Other,IKE,MAIL,SOCKS,DoH,SLP,SSL with ESNI,ISATAP,Stratum,SSL with ECH + + - function: GEOIP_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ ] + parameters: + kb_name: cn_ip_location + option: IP_TO_OBJECT + geolocation_field_mapping: + COUNTRY: client_country_region + PROVINCE: client_super_admin_area + CITY: client_admin_area + LONGITUDE: client_longitude + LATITUDE: client_latitude + ISP: client_isp + + - function: GEOIP_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ ] + parameters: + kb_name: cn_ip_location + option: IP_TO_OBJECT + geolocation_field_mapping: + COUNTRY: server_country_region + PROVINCE: server_super_admin_area + CITY: server_admin_area + LONGITUDE: server_longitude + LATITUDE: server_latitude + ISP: server_isp + + - function: ASN_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_asn ] + parameters: + option: IP_TO_ASN + kb_name: cn_ip_asn + + - function: ASN_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_asn ] + parameters: + option: IP_TO_ASN + kb_name: cn_ip_asn + + - function: CN_IDC_RENTER_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_idc_renter ] + parameters: + kb_name: cn_idc_renter + + - function: CN_IDC_RENTER_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_idc_renter ] + parameters: + kb_name: cn_idc_renter + + - function: CN_LINK_DIRECTION_LOOKUP + lookup_fields: [ in_link_id ] + output_fields: [ in_link_direction ] + parameters: + kb_name: cn_link_direction + + - function: CN_LINK_DIRECTION_LOOKUP + lookup_fields: [ out_link_id ] + output_fields: [ out_link_direction ] + parameters: + kb_name: cn_link_direction + + - function: CN_ICP_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_icp_company_name ] + parameters: + kb_name: cn_fqdn_icp + + - function: CN_FQDN_WHOIS_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_whois_org ] + parameters: + kb_name: cn_fqdn_whois + + - function: CN_DNS_SERVER_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_dns_server ] + parameters: + kb_name: cn_dns_server + + - function: CN_APP_CATEGORY_LOOKUP + lookup_fields: [ app ] + parameters: + kb_name: cn_app_category + field_mapping: + CATEGORY: app_category + SUBCATEGORY: app_subcategory + COMPANY: app_company + COMPANY_CATEGORY: app_company_category + + - function: CN_IP_ZONE_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_zone ] + parameters: + kb_name: cn_internal_ip + + - function: CN_IP_ZONE_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_zone ] + parameters: + kb_name: cn_internal_ip + + - function: CN_VPN_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_vpn_service_name ] + parameters: + kb_name: cn_vpn_learning_ip + option: IP_TO_VPN + + - function: CN_VPN_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_vpn_service_name ] + parameters: + kb_name: cn_vpn_learning_domain + option: DOMAIN_TO_VPN + + - function: CN_IOC_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_malware ] + parameters: + kb_name: cn_ioc_malware + option: IP_TO_MALWARE + + - function: CN_IOC_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_malware ] + parameters: + kb_name: cn_ioc_malware + option: DOMAIN_TO_MALWARE + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_ip_tags ] + parameters: + kb_name: cn_ip_tag_user_define + option: IP_TO_TAG + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_ip_tags ] + parameters: + kb_name: cn_ip_tag_user_define + option: IP_TO_TAG + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_tags ] + parameters: + kb_name: cn_domain_tag_user_define + option: DOMAIN_TO_TAG + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ app ] + output_fields: [ app_tags ] + parameters: + kb_name: cn_app_tag_user_define + option: APP_TO_TAG + + - function: CN_FIELDS_MERGE + lookup_fields: [ client_idc_renter,client_ip_tags ] + output_fields: [ client_ip_tags ] + + - function: CN_FIELDS_MERGE + lookup_fields: [ server_idc_renter,server_dns_server,server_node_type,server_malware,server_vpn_service_name,server_ip_tags ] + output_fields: [ server_ip_tags ] + + - function: CN_FIELDS_MERGE + lookup_fields: [ domain_node_type,domain_malware,domain_vpn_service_name,domain_tags ] + output_fields: [ domain_tags ] + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ client_ip_tags ] + output_fields: [ client_ip_tags ] + parameters: + prefix: ip. + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ server_ip_tags ] + output_fields: [ server_ip_tags ] + parameters: + prefix: ip. + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ domain_tags ] + output_fields: [ domain_tags ] + parameters: + prefix: domain. + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ app_tags ] + output_fields: [ app_tags ] + parameters: + prefix: app. + +postprocessing_pipelines: + remove_field_processor: # [object] Processing Pipeline + type: com.geedgenetworks.core.processor.projection.ProjectionProcessorImpl + remove_fields: [ log_id, device_tag, dup_traffic_flag ] + +sinks: + kafka_sink_a: + type: kafka + properties: + topic: test + kafka.bootstrap.servers: 192.168.44.55:9092 + kafka.retries: 0 + kafka.linger.ms: 10 + kafka.request.timeout.ms: 30000 + kafka.batch.size: 262144 + kafka.buffer.memory: 134217728 + kafka.max.request.size: 10485760 + kafka.compression.type: snappy + kafka.security.protocol: + kafka.ssl.keystore.location: + kafka.ssl.keystore.password: + kafka.ssl.truststore.location: + kafka.ssl.truststore.password: + kafka.ssl.key.password: + kafka.sasl.mechanism: + kafka.sasl.jaas.config: + format: json + + print_sink: + type: print + properties: + format: json + +application: # [object] Application Configuration + env: # [object] Environment Variables + name: groot-stream-job # [string] Job Name + parallelism: 3 # [number] Job-Level Parallelism + pipeline: + object-reuse: true # [boolean] Object Reuse, default is false + topology: # [array of object] Node List. It will be used build data flow for job dag graph. + - name: kafka_source # [string] Node Name, must be unique. It will be used as the name of the corresponding Flink operator. eg. kafka_source the processor type as SOURCE. + #parallelism: 1 # [number] Operator-Level Parallelism. + downstream: [ session_record_processor ] # [array of string] Downstream Node Name List. + - name: session_record_processor + downstream: [ remove_field_processor ] + - name: remove_field_processor + downstream: [ print_sink ] + - name: kafka_sink_a + downstream: [ ] + - name: print_sink + downstream: [ ] diff --git a/config/cn/cn_grootstream_job_template.yaml b/config/cn/cn_grootstream_job_template.yaml new file mode 100644 index 0000000..7c448f6 --- /dev/null +++ b/config/cn/cn_grootstream_job_template.yaml @@ -0,0 +1,339 @@ +sources: + kafka_source: + type: kafka + # fields: # [array of object] Field List, if not set, all fields(Map<String, Object>) will be output. + properties: # [object] Source Properties + topic: SESSION-RECORD + kafka.bootstrap.servers: 192.168.44.11:9094,192.168.44.13:9094,192.168.44.14:9094,192.168.44.15:9094,192.168.44.16:9094 + kafka.session.timeout.ms: 60000 + kafka.max.poll.records: 3000 + kafka.max.partition.fetch.bytes: 31457280 + kafka.security.protocol: SASL_PLAINTEXT + kafka.ssl.keystore.location: + kafka.ssl.keystore.password: + kafka.ssl.truststore.location: + kafka.ssl.truststore.password: + kafka.ssl.key.password: + kafka.sasl.mechanism: PLAIN + kafka.sasl.jaas.config: org.apache.kafka.common.security.plain.PlainLoginModule required username="admin" password="galaxy2019"; + kafka.buffer.memory: + kafka.group.id: local-test + kafka.auto.offset.reset: latest + kafka.max.request.size: + kafka.compression.type: none + format: json # [string] Data Format, default is json + +processing_pipelines: + session_record_processor: # [object] Processing Pipeline + type: com.geedgenetworks.core.processor.projection.ProjectionProcessorImpl + remove_fields: + output_fields: + functions: # [array of object] Function List + - function: SNOWFLAKE_ID + lookup_fields: [ '' ] + output_fields: [ log_id ] + filter: + parameters: + data_center_id_num: 1 + + - function: UNIX_TIMESTAMP_CONVERTER + lookup_fields: [ __timestamp ] + output_fields: [ recv_time ] + parameters: + precision: seconds + + - function: EVAL + output_fields: [ ingestion_time ] + parameters: + value_expression: recv_time + + - function: EVAL + output_fields: [ domain ] + parameters: + value_expression: server_fqdn + + - function: EVAL + output_fields: [ domain_sld ] + parameters: + value_expression: server_domain + + - function: CN_L7_PROTOCOL_AND_APP_EXTRACT + parameters: + decoded_path_field_name: decoded_path + app_transition_field_name: app_transition + l7_protocol_field_name: l7_protocol + app_field_name: app + l7_protocol: DHCP,DNS,FTP,GRE,GTP,HTTP,HTTPS,ICMP,IMAP,IMAPS,IPSEC,ISAKMP,XMPP,L2TP,LDAP,MMS,NETBIOS,NETFLOW,NTP,POP3,POP3S,RDP,PPTP,RADIUS,RTCP,RTP,RTSP,SIP,SMB,SMTP,SMTPS,SNMP,SSDP,SSH,SSL,STUN,TELNET,TFTP,OPENVPN,RTMP,TEREDO,FTPS,DTLS,SPDY,BJNP,QUIC,MDNS,Unknown TCP,Unknown UDP,Unknown Other,IKE,MAIL,SOCKS,DoH,SLP,SSL with ESNI,ISATAP,Stratum,SSL with ECH + + - function: GEOIP_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ ] + parameters: + kb_name: cn_ip_location + option: IP_TO_OBJECT + geolocation_field_mapping: + COUNTRY: client_country_region + PROVINCE: client_super_admin_area + CITY: client_admin_area + LONGITUDE: client_longitude + LATITUDE: client_latitude + ISP: client_isp + + - function: GEOIP_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ ] + parameters: + kb_name: cn_ip_location + option: IP_TO_OBJECT + geolocation_field_mapping: + COUNTRY: server_country_region + PROVINCE: server_super_admin_area + CITY: server_admin_area + LONGITUDE: server_longitude + LATITUDE: server_latitude + ISP: server_isp + + - function: ASN_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_asn ] + parameters: + option: IP_TO_ASN + kb_name: cn_ip_asn + + - function: ASN_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_asn ] + parameters: + option: IP_TO_ASN + kb_name: cn_ip_asn + + - function: CN_IDC_RENTER_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_idc_renter ] + parameters: + kb_name: cn_idc_renter + + - function: CN_IDC_RENTER_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_idc_renter ] + parameters: + kb_name: cn_idc_renter + + - function: CN_LINK_DIRECTION_LOOKUP + lookup_fields: [ in_link_id ] + output_fields: [ in_link_direction ] + parameters: + kb_name: cn_link_direction + + - function: CN_LINK_DIRECTION_LOOKUP + lookup_fields: [ out_link_id ] + output_fields: [ out_link_direction ] + parameters: + kb_name: cn_link_direction + + - function: CN_FQDN_CATEGORY_LOOKUP + lookup_fields: [ domain ] + parameters: + kb_name: cn_fqdn_category + field_mapping: + NAME: domain_category_name + GROUP: domain_category_group + REPUTATION_LEVEL: domain_reputation_level + + - function: CN_ICP_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_icp_company_name ] + parameters: + kb_name: cn_fqdn_icp + + - function: CN_FQDN_WHOIS_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_whois_org ] + parameters: + kb_name: cn_fqdn_whois + + - function: CN_DNS_SERVER_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_dns_server ] + parameters: + kb_name: cn_dns_server + + - function: CN_APP_CATEGORY_LOOKUP + lookup_fields: [ app ] + parameters: + kb_name: cn_app_category + field_mapping: + CATEGORY: app_category + SUBCATEGORY: app_subcategory + COMPANY: app_company + COMPANY_CATEGORY: app_company_category + + - function: CN_IP_ZONE_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_zone ] + parameters: + kb_name: cn_internal_ip + + - function: CN_IP_ZONE_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_zone ] + parameters: + kb_name: cn_internal_ip + + - function: CN_VPN_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_vpn_service_name ] + parameters: + kb_name: cn_vpn_learning_ip + option: IP_TO_VPN + + - function: CN_VPN_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_vpn_service_name ] + parameters: + kb_name: cn_vpn_learning_domain + option: DOMAIN_TO_VPN + + - function: CN_ANONYMITY_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_node_type ] + parameters: + kb_name: cn_ioc_darkweb + option: IP_TO_NODE_TYPE + + - function: CN_ANONYMITY_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_node_type ] + parameters: + kb_name: cn_ioc_darkweb + option: DOMAIN_TO_NODE_TYPE + + - function: CN_IOC_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_malware ] + parameters: + kb_name: cn_ioc_malware + option: IP_TO_MALWARE + + - function: CN_IOC_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_malware ] + parameters: + kb_name: cn_ioc_malware + option: DOMAIN_TO_MALWARE + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ client_ip ] + output_fields: [ client_ip_tags ] + parameters: + kb_name: cn_ip_tag_user_define + option: IP_TO_TAG + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ server_ip ] + output_fields: [ server_ip_tags ] + parameters: + kb_name: cn_ip_tag_user_define + option: IP_TO_TAG + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ domain ] + output_fields: [ domain_tags ] + parameters: + kb_name: cn_domain_tag_user_define + option: DOMAIN_TO_TAG + + - function: CN_USER_DEFINE_TAG_LOOKUP + lookup_fields: [ app ] + output_fields: [ app_tags ] + parameters: + kb_name: cn_app_tag_user_define + option: APP_TO_TAG + + - function: CN_FIELDS_MERGE + lookup_fields: [ client_idc_renter,client_ip_tags ] + output_fields: [ client_ip_tags ] + + - function: CN_FIELDS_MERGE + lookup_fields: [ server_idc_renter,server_dns_server,server_node_type,server_malware,server_vpn_service_name,server_ip_tags ] + output_fields: [ server_ip_tags ] + + - function: CN_FIELDS_MERGE + lookup_fields: [ domain_node_type,domain_malware,domain_vpn_service_name,domain_tags ] + output_fields: [ domain_tags ] + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ client_ip_tags ] + output_fields: [ client_ip_tags ] + parameters: + prefix: ip. + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ server_ip_tags ] + output_fields: [ server_ip_tags ] + parameters: + prefix: ip. + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ domain_tags ] + output_fields: [ domain_tags ] + parameters: + prefix: domain. + + - function: CN_ARRAY_ELEMENTS_PREPEND + lookup_fields: [ app_tags ] + output_fields: [ app_tags ] + parameters: + prefix: app. + +postprocessing_pipelines: + remove_field_processor: # [object] Processing Pipeline + type: com.geedgenetworks.core.processor.projection.ProjectionProcessorImpl + remove_fields: [ log_id, device_tag, dup_traffic_flag ] + +sinks: + kafka_sink_a: + type: kafka + properties: + topic: test + kafka.bootstrap.servers: 192.168.44.55:9092 + kafka.retries: 0 + kafka.linger.ms: 10 + kafka.request.timeout.ms: 30000 + kafka.batch.size: 262144 + kafka.buffer.memory: 134217728 + kafka.max.request.size: 10485760 + kafka.compression.type: snappy + kafka.security.protocol: + kafka.ssl.keystore.location: + kafka.ssl.keystore.password: + kafka.ssl.truststore.location: + kafka.ssl.truststore.password: + kafka.ssl.key.password: + kafka.sasl.mechanism: + kafka.sasl.jaas.config: + format: json + + print_sink: + type: print + properties: + format: json + +application: # [object] Application Configuration + env: # [object] Environment Variables + name: groot-stream-job # [string] Job Name + parallelism: 3 # [number] Job-Level Parallelism + pipeline: + object-reuse: true # [boolean] Object Reuse, default is false + topology: # [array of object] Node List. It will be used build data flow for job dag graph. + - name: kafka_source # [string] Node Name, must be unique. It will be used as the name of the corresponding Flink operator. eg. kafka_source the processor type as SOURCE. + #parallelism: 1 # [number] Operator-Level Parallelism. + downstream: [ session_record_processor ] # [array of string] Downstream Node Name List. + - name: session_record_processor + downstream: [ remove_field_processor ] + - name: remove_field_processor + downstream: [ print_sink ] + - name: kafka_sink_a + downstream: [ ] + - name: print_sink + downstream: [ ] diff --git a/config/cn/grootstream.yaml b/config/cn/grootstream.yaml new file mode 100644 index 0000000..558030c --- /dev/null +++ b/config/cn/grootstream.yaml @@ -0,0 +1,108 @@ +grootstream: + knowledge_base: + - name: cn_ip_location + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 1 + + - name: cn_ip_asn + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 2 + + - name: cn_idc_renter + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 11 + + - name: cn_link_direction + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 13 + + - name: cn_fqdn_category + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 5 + + - name: cn_fqdn_icp + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 4 + + - name: cn_fqdn_whois + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 6 + + - name: cn_dns_server + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 3 + + - name: cn_app_category + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 9 + + - name: cn_internal_ip + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 12 + + - name: cn_vpn_learning_ip + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 15 + + - name: cn_vpn_learning_domain + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 14 + + - name: cn_ioc_darkweb + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 8 + + - name: cn_ioc_malware + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base + files: + - 7 + + - name: cn_ip_tag_user_define + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base?category=cn_ip_tag_user_defined + + - name: cn_domain_tag_user_define + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base?category=cn_domain_tag_user_defined + + - name: cn_app_tag_user_define + fs_type: http + fs_path: http://192.168.44.55:9999/v1/knowledge_base?category=cn_app_tag_user_defined + + - name: cn_rule + fs_type: http + fs_path: http://192.168.44.54:8090 + properties: + token: 1a653ea0-d39b-4246-94b0-1ba95db4b6a7 + + properties: + hos.path: http://192.168.44.12:8089 + hos.bucket.name.traffic_file: traffic_file_bucket + hos.bucket.name.troubleshooting_file: troubleshooting_file_bucket
\ No newline at end of file diff --git a/config/cn/udf.plugins b/config/cn/udf.plugins new file mode 100644 index 0000000..22804f6 --- /dev/null +++ b/config/cn/udf.plugins @@ -0,0 +1,30 @@ +com.geedgenetworks.core.udf.AsnLookup +com.geedgenetworks.core.udf.CurrentUnixTimestamp +com.geedgenetworks.core.udf.DecodeBase64 +com.geedgenetworks.core.udf.Domain +com.geedgenetworks.core.udf.Drop +com.geedgenetworks.core.udf.Eval +com.geedgenetworks.core.udf.FromUnixTimestamp +com.geedgenetworks.core.udf.GenerateStringArray +com.geedgenetworks.core.udf.GeoIpLookup +com.geedgenetworks.core.udf.JsonExtract +com.geedgenetworks.core.udf.PathCombine +com.geedgenetworks.core.udf.Rename +com.geedgenetworks.core.udf.SnowflakeId +com.geedgenetworks.core.udf.StringJoiner +com.geedgenetworks.core.udf.UnixTimestampConverter +com.geedgenetworks.core.udf.cn.L7ProtocolAndAppExtract +com.geedgenetworks.core.udf.cn.IdcRenterLookup +com.geedgenetworks.core.udf.cn.LinkDirectionLookup +com.geedgenetworks.core.udf.cn.FqdnCategoryLookup +com.geedgenetworks.core.udf.cn.IcpLookup +com.geedgenetworks.core.udf.cn.FqdnWhoisLookup +com.geedgenetworks.core.udf.cn.DnsServerInfoLookup +com.geedgenetworks.core.udf.cn.AppCategoryLookup +com.geedgenetworks.core.udf.cn.IpZoneLookup +com.geedgenetworks.core.udf.cn.VpnLookup +com.geedgenetworks.core.udf.cn.AnonymityLookup +com.geedgenetworks.core.udf.cn.IocLookup +com.geedgenetworks.core.udf.cn.UserDefineTagLookup +com.geedgenetworks.core.udf.cn.FieldsMerge +com.geedgenetworks.core.udf.cn.ArrayElementsPrepend diff --git a/groot-core/pom.xml b/groot-core/pom.xml index c14b5fb..f19e4b1 100644 --- a/groot-core/pom.xml +++ b/groot-core/pom.xml @@ -14,6 +14,30 @@ <dependencies> <dependency> + <groupId>org.mockito</groupId> + <artifactId>mockito-core</artifactId> + <version>4.0.0</version> + <scope>test</scope> + </dependency> + + <dependency> + <groupId>org.mockito</groupId> + <artifactId>mockito-inline</artifactId> + <version>4.0.0</version> + <scope>test</scope> + </dependency> + + <dependency> + <groupId>com.github.seancfoley</groupId> + <artifactId>ipaddress</artifactId> + </dependency> + + <dependency> + <groupId>com.opencsv</groupId> + <artifactId>opencsv</artifactId> + </dependency> + + <dependency> <groupId>com.geedgenetworks</groupId> <artifactId>http-client-shaded</artifactId> <version>${project.version}</version> diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AbstractKnowledgeUDF.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AbstractKnowledgeUDF.java new file mode 100644 index 0000000..bf21b62 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AbstractKnowledgeUDF.java @@ -0,0 +1,48 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Constants; +import com.geedgenetworks.common.config.CommonConfig; +import com.geedgenetworks.common.config.KnowledgeBaseConfig; +import com.geedgenetworks.common.udf.UDF; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.apache.flink.configuration.Configuration; + +import java.util.List; +import java.util.stream.Collectors; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 9:33 + */ +public abstract class AbstractKnowledgeUDF implements UDF { + + protected String lookupFieldName; + + protected String outputFieldName; + + protected String cnInternalFieldNamePrefix = "cn_internal_"; + + protected List<KnowledgeBaseConfig> knowledgeBaseConfigs; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + String kbName = udfContext.getParameters().get("kb_name").toString(); + + Configuration configuration = (Configuration) runtimeContext + .getExecutionConfig().getGlobalJobParameters(); + CommonConfig commonConfig = com.alibaba.fastjson.JSON.parseObject(configuration.toMap().get(Constants.SYSPROP_GROOTSTREAM_CONFIG), CommonConfig.class); + + knowledgeBaseConfigs = commonConfig.getKnowledgeBaseConfig().stream().filter(knowledgeBaseConfig -> knowledgeBaseConfig.getName().equals(kbName)).collect(Collectors.toList()); + + registerKnowledges(); + + this.lookupFieldName = udfContext.getLookup_fields().get(0); + if (udfContext.getOutput_fields() != null && udfContext.getOutput_fields().size() > 0) { + this.outputFieldName = udfContext.getOutput_fields().get(0); + } + } + + protected abstract void registerKnowledges(); +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AbstractKnowledgeWithRuleUDF.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AbstractKnowledgeWithRuleUDF.java new file mode 100644 index 0000000..d0786d6 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AbstractKnowledgeWithRuleUDF.java @@ -0,0 +1,74 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Constants; +import com.geedgenetworks.common.config.CommonConfig; +import com.geedgenetworks.common.config.KnowledgeBaseConfig; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.apache.flink.configuration.Configuration; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 18:15 + */ +public abstract class AbstractKnowledgeWithRuleUDF extends AbstractKnowledgeUDF { + + protected List<KnowledgeBaseConfig> ruleConfigs; + + protected String internalRuleIdListFieldName = cnInternalFieldNamePrefix + "rule_id_list"; + + protected String internalIocTypeListFieldName = cnInternalFieldNamePrefix + "ioc_type_list"; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + Configuration configuration = (Configuration) runtimeContext + .getExecutionConfig().getGlobalJobParameters(); + + CommonConfig commonConfig = com.alibaba.fastjson.JSON.parseObject(configuration.toMap().get(Constants.SYSPROP_GROOTSTREAM_CONFIG), CommonConfig.class); + + ruleConfigs = commonConfig.getKnowledgeBaseConfig().stream().filter(knowledgeBaseConfig -> knowledgeBaseConfig.getName().equals("cn_rule")).collect(Collectors.toList()); + + super.open(runtimeContext, udfContext); + } + + protected enum IocType { + IP("ip"), + DOMAIN("domain"), + APP("app"); + + private String type; + + IocType(String type) { + this.type = type; + } + + public String getType() { + return type; + } + } + + protected static final class RuleMetadata { + + List<Long> ruleIds = new ArrayList<>(); + List<String> iocTypes = new ArrayList<>(); + + public void addRule(Long ruleId, String iocType) { + ruleIds.add(ruleId); + iocTypes.add(iocType); + } + + public List<Long> getRuleIds() { + return Collections.unmodifiableList(ruleIds); + } + + public List<String> getIocTypes() { + return Collections.unmodifiableList(iocTypes); + } + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AnonymityLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AnonymityLookup.java new file mode 100644 index 0000000..0e27355 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AnonymityLookup.java @@ -0,0 +1,86 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.IocDarkwebKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.RuleKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 14:25 + */ +public class AnonymityLookup extends AbstractKnowledgeWithRuleUDF { + + private static final Logger logger = LoggerFactory.getLogger(AnonymityLookup.class); + + private String option; + + private IocDarkwebKnowledgeBaseHandler knowledgeBaseHandler; + + private RuleKnowledgeBaseHandler ruleKnowledgeBaseHandler; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + option = udfContext.getParameters().get("option").toString(); + super.open(runtimeContext, udfContext); + } + + @Override + @SuppressWarnings("unchecked") + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + String lookupValue = event.getExtractedFields().get(lookupFieldName).toString(); + RuleMetadata ruleMetadata = new RuleMetadata(); + switch (option) { + case "IP_TO_NODE_TYPE": + String ipNodeType = knowledgeBaseHandler.lookupByIp(lookupValue); + event.getExtractedFields().put(outputFieldName, ipNodeType); + RuleKnowledgeBaseHandler.Rule ipRule = ruleKnowledgeBaseHandler.lookupByName(ipNodeType); + if (ipRule != null) { + ruleMetadata.addRule(ipRule.getRuleId(), IocType.IP.getType()); + } + break; + case "DOMAIN_TO_NODE_TYPE": + String domainNodeType = knowledgeBaseHandler.lookupByDomain(lookupValue); + event.getExtractedFields().put(outputFieldName, domainNodeType); + RuleKnowledgeBaseHandler.Rule domainRule = ruleKnowledgeBaseHandler.lookupByName(domainNodeType); + if (domainRule != null) { + ruleMetadata.addRule(domainRule.getRuleId(), IocType.DOMAIN.getType()); + } + break; + default: + logger.error("unknown option: " + option); + break; + } + (((ArrayList<Long>) (event.getExtractedFields().computeIfAbsent(internalRuleIdListFieldName, k -> new ArrayList<Long>())))).addAll(ruleMetadata.getRuleIds()); + (((ArrayList<String>) (event.getExtractedFields().computeIfAbsent(internalIocTypeListFieldName, k -> new ArrayList<String>())))).addAll(ruleMetadata.getIocTypes()); + } + return event; + } + + @Override + public String functionName() { + return "CN_ANONYMITY_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + ruleKnowledgeBaseHandler = RuleKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(ruleKnowledgeBaseHandler, ruleConfigs.get(0)); + + knowledgeBaseHandler = IocDarkwebKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AppCategoryLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AppCategoryLookup.java new file mode 100644 index 0000000..feae361 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/AppCategoryLookup.java @@ -0,0 +1,76 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.AppCategoryKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 14:14 + */ +public class AppCategoryLookup extends AbstractKnowledgeUDF { + + private static final Logger logger = LoggerFactory.getLogger(AppCategoryLookup.class); + + private AppCategoryKnowledgeBaseHandler knowledgeBaseHandler; + + private Map<String, String> fieldMapping; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + super.open(runtimeContext, udfContext); + fieldMapping = (Map<String, String>) udfContext.getParameters().get("field_mapping"); + } + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + AppCategoryKnowledgeBaseHandler.AppCategory appCategory = knowledgeBaseHandler.lookup(event.getExtractedFields().get(lookupFieldName).toString()); + if (appCategory != null) { + fieldMapping.forEach((key, value) -> { + switch (key) { + case "CATEGORY": + event.getExtractedFields().put(value, appCategory.getCategory()); + break; + case "SUBCATEGORY": + event.getExtractedFields().put(value, appCategory.getSubCategory()); + break; + case "COMPANY": + event.getExtractedFields().put(value, appCategory.getAppCompany()); + break; + case "COMPANY_CATEGORY": + event.getExtractedFields().put(value, appCategory.getAppCompanyCategory()); + break; + default: + logger.error("unknown field name :" + key); + break; + } + }); + } + } + return event; + } + + @Override + public String functionName() { + return "CN_APP_CATEGORY_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = AppCategoryKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/ArrayElementsPrepend.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/ArrayElementsPrepend.java new file mode 100644 index 0000000..27bd228 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/ArrayElementsPrepend.java @@ -0,0 +1,50 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDF; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; + +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Collectors; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/23 11:14 + */ +public class ArrayElementsPrepend implements UDF { + + private String prefix; + + private String lookupFieldName; + private String outputFieldName; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + this.prefix = udfContext.getParameters().get("prefix").toString(); + this.lookupFieldName = udfContext.getLookup_fields().get(0); + this.outputFieldName = udfContext.getOutput_fields().get(0); + } + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName) instanceof ArrayList) { + List<String> list = ((ArrayList<?>) (event.getExtractedFields().get(lookupFieldName))).stream().map(s -> prefix + s).collect(Collectors.toList()); + event.getExtractedFields().put(outputFieldName, list); + return event; + } + return event; + } + + @Override + public String functionName() { + return "CN_ARRAY_ELEMENTS_PREPEND"; + } + + @Override + public void close() { + + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/DnsServerInfoLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/DnsServerInfoLookup.java new file mode 100644 index 0000000..3812bc6 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/DnsServerInfoLookup.java @@ -0,0 +1,44 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.DnsServerInfoKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/18 16:21 + */ +public class DnsServerInfoLookup extends AbstractKnowledgeUDF { + + private DnsServerInfoKnowledgeBaseHandler knowledgeBaseHandler; + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + event.getExtractedFields().put( + outputFieldName, + knowledgeBaseHandler.lookup( + event.getExtractedFields().get(lookupFieldName).toString() + ) + ); + } + return event; + } + + @Override + public String functionName() { + return "CN_DNS_SERVER_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = DnsServerInfoKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FieldsMerge.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FieldsMerge.java new file mode 100644 index 0000000..ae53321 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FieldsMerge.java @@ -0,0 +1,53 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDF; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; + +import java.util.ArrayList; +import java.util.List; +import java.util.stream.Collectors; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/23 11:20 + */ +public class FieldsMerge implements UDF { + + private List<String> lookupFieldNames; + private String outputFieldName; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + this.lookupFieldNames = udfContext.getLookup_fields(); + this.outputFieldName = udfContext.getOutput_fields().get(0); + } + + @Override + public Event evaluate(Event event) { + List<String> list = new ArrayList<>(); + lookupFieldNames.forEach(fieldName -> { + Object value = event.getExtractedFields().get(fieldName); + if (value instanceof ArrayList) { + list.addAll(((ArrayList<?>) value).stream().map(Object::toString).collect(Collectors.toList())); + } + if (value instanceof String) { + list.add((String) value); + } + }); + event.getExtractedFields().put(outputFieldName, list); + return event; + } + + @Override + public String functionName() { + return "CN_FIELDS_MERGE"; + } + + @Override + public void close() { + + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FqdnCategoryLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FqdnCategoryLookup.java new file mode 100644 index 0000000..4aa007c --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FqdnCategoryLookup.java @@ -0,0 +1,73 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.FqdnCategoryKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 10:22 + */ +public class FqdnCategoryLookup extends AbstractKnowledgeUDF { + + private static final Logger logger = LoggerFactory.getLogger(FqdnCategoryLookup.class); + + private Map<String, String> fieldMapping; + + private FqdnCategoryKnowledgeBaseHandler knowledgeBaseHandler; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + super.open(runtimeContext, udfContext); + fieldMapping = (Map<String, String>) udfContext.getParameters().get("field_mapping"); + } + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + FqdnCategoryKnowledgeBaseHandler.FqdnCategory fqdnCategory = knowledgeBaseHandler.lookup(event.getExtractedFields().get(lookupFieldName).toString()); + if (fqdnCategory != null) { + fieldMapping.forEach((key, value) -> { + switch (key) { + case "NAME": + event.getExtractedFields().put(value, fqdnCategory.getCategoryName()); + break; + case "GROUP": + event.getExtractedFields().put(value, fqdnCategory.getCategoryGroup()); + break; + case "REPUTATION_LEVEL": + event.getExtractedFields().put(value, fqdnCategory.getReputationLevel()); + break; + default: + logger.error("unknown field name :" + key); + break; + } + }); + } + } + return event; + } + + @Override + public String functionName() { + return "CN_FQDN_CATEGORY_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = FqdnCategoryKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FqdnWhoisLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FqdnWhoisLookup.java new file mode 100644 index 0000000..28b5282 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/FqdnWhoisLookup.java @@ -0,0 +1,44 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.FqdnWhoisKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 13:59 + */ +public class FqdnWhoisLookup extends AbstractKnowledgeUDF { + + private FqdnWhoisKnowledgeBaseHandler knowledgeBaseHandler; + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + event.getExtractedFields().put( + outputFieldName, + knowledgeBaseHandler.lookup( + event.getExtractedFields().get(lookupFieldName).toString() + ) + ); + } + return event; + } + + @Override + public String functionName() { + return "CN_FQDN_WHOIS_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = FqdnWhoisKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IcpLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IcpLookup.java new file mode 100644 index 0000000..395e8ff --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IcpLookup.java @@ -0,0 +1,44 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.FqdnIcpKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 10:27 + */ +public class IcpLookup extends AbstractKnowledgeUDF { + + private FqdnIcpKnowledgeBaseHandler knowledgeBaseHandler; + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + event.getExtractedFields().put( + outputFieldName, + knowledgeBaseHandler.lookup( + event.getExtractedFields().get(lookupFieldName).toString() + ) + ); + } + return event; + } + + @Override + public String functionName() { + return "CN_ICP_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = FqdnIcpKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IdcRenterLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IdcRenterLookup.java new file mode 100644 index 0000000..dfb63d5 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IdcRenterLookup.java @@ -0,0 +1,44 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.IdcRenterKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/18 11:45 + */ +public class IdcRenterLookup extends AbstractKnowledgeUDF { + + private IdcRenterKnowledgeBaseHandler knowledgeBaseHandler; + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + event.getExtractedFields().put( + outputFieldName, + knowledgeBaseHandler.lookup( + event.getExtractedFields().get(lookupFieldName).toString() + ) + ); + } + return event; + } + + @Override + public String functionName() { + return "CN_IDC_RENTER_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = IdcRenterKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IocLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IocLookup.java new file mode 100644 index 0000000..03dbace --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IocLookup.java @@ -0,0 +1,94 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.IocMalwareKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.RuleKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 14:45 + */ +public class IocLookup extends AbstractKnowledgeWithRuleUDF { + + private static final Logger logger = LoggerFactory.getLogger(IocLookup.class); + + private String option; + + private IocMalwareKnowledgeBaseHandler knowledgeBaseHandler; + + private RuleKnowledgeBaseHandler ruleKnowledgeBaseHandler; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + option = udfContext.getParameters().get("option").toString(); + super.open(runtimeContext, udfContext); + } + + @Override + @SuppressWarnings("unchecked") + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + String lookupValue = event.getExtractedFields().get(lookupFieldName).toString(); + RuleMetadata ruleMetadata = new RuleMetadata(); + switch (option) { + case "IP_TO_MALWARE": + String ipMalware = knowledgeBaseHandler.lookupByIp(lookupValue); + event.getExtractedFields().put(outputFieldName, ipMalware); + RuleKnowledgeBaseHandler.Rule ipRule = ruleKnowledgeBaseHandler.lookupByName(ipMalware); + if (ipRule != null) { + ruleMetadata.addRule(ipRule.getRuleId(), IocType.IP.getType()); + } + break; + case "DOMAIN_TO_MALWARE": + String domainMalware = knowledgeBaseHandler.lookupByDomain(lookupValue); + event.getExtractedFields().put(outputFieldName, domainMalware); + RuleKnowledgeBaseHandler.Rule domainRule = ruleKnowledgeBaseHandler.lookupByName(domainMalware); + if (domainRule != null) { + ruleMetadata.addRule(domainRule.getRuleId(), IocType.DOMAIN.getType()); + } + break; + case "HTTP_URL_TO_MALWARE": + String urlMalware = knowledgeBaseHandler.lookupByUrl(lookupValue); + event.getExtractedFields().put(outputFieldName, urlMalware); + RuleKnowledgeBaseHandler.Rule urlRule = ruleKnowledgeBaseHandler.lookupByName(urlMalware); + if (urlRule != null) { + ruleMetadata.addRule(urlRule.getRuleId(), IocType.DOMAIN.getType()); + } + break; + default: + logger.error("unknown option: " + option); + break; + } + (((ArrayList<Long>) (event.getExtractedFields().computeIfAbsent(internalRuleIdListFieldName, k -> new ArrayList<Long>())))).addAll(ruleMetadata.getRuleIds()); + (((ArrayList<String>) (event.getExtractedFields().computeIfAbsent(internalIocTypeListFieldName, k -> new ArrayList<String>())))).addAll(ruleMetadata.getIocTypes()); + } + return event; + } + + @Override + public String functionName() { + return "CN_IOC_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + ruleKnowledgeBaseHandler = RuleKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(ruleKnowledgeBaseHandler, ruleConfigs.get(0)); + + knowledgeBaseHandler = IocMalwareKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IpZoneLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IpZoneLookup.java new file mode 100644 index 0000000..c174deb --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IpZoneLookup.java @@ -0,0 +1,41 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.InternalIpKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 10:19 + */ +public class IpZoneLookup extends AbstractKnowledgeUDF { + + private InternalIpKnowledgeBaseHandler knowledgeBaseHandler; + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + String ip = event.getExtractedFields().get(lookupFieldName).toString(); + String zone = knowledgeBaseHandler.isInternal(ip) ? "internal" : "external"; + event.getExtractedFields().put(outputFieldName, zone); + } + return event; + } + + @Override + public String functionName() { + return "CN_IP_ZONE_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = InternalIpKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/L7ProtocolAndAppExtract.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/L7ProtocolAndAppExtract.java new file mode 100644 index 0000000..897e285 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/L7ProtocolAndAppExtract.java @@ -0,0 +1,98 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDF; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/23 10:59 + */ +public class L7ProtocolAndAppExtract implements UDF { + + private static final Logger logger = LoggerFactory.getLogger(L7ProtocolAndAppExtract.class); + + private String decodedPathFieldName; + private String appTransitionFieldName; + private String l7ProtocolFieldName; + private String appFieldName; + + private HashMap<String, String> l7ProtocolMap = new HashMap<>(); + + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + decodedPathFieldName = udfContext.getParameters().get("decoded_path_field_name").toString(); + appTransitionFieldName = udfContext.getParameters().get("app_transition_field_name").toString(); + l7ProtocolFieldName = udfContext.getParameters().get("l7_protocol_field_name").toString(); + appFieldName = udfContext.getParameters().get("app_field_name").toString(); + String l7Protocol = udfContext.getParameters().get("l7_protocol").toString(); + l7ProtocolMap = Arrays.stream(l7Protocol.split(",")).collect(Collectors.toMap(String::toLowerCase, Function.identity(), (existing, replacement) -> existing, HashMap::new)); + } + + @Override + public Event evaluate(Event event) { + try { + String protocolPath = event.getExtractedFields().computeIfAbsent(decodedPathFieldName, k -> "").toString(); + String appFullPath = event.getExtractedFields().computeIfAbsent(appTransitionFieldName, k -> "").toString(); + String fullPath = ""; + int appIndex = 0; + if (!"".equals(protocolPath)) { + appIndex = protocolPath.split("\\.").length; + fullPath = protocolPath; + } + if (!"".equals(appFullPath)) { + fullPath = fullPath + "." + appFullPath; + } + String l7Protocol = "UNCATEGORIZED"; + String app = ""; + if (!"".equals(fullPath)) { + fullPath = fullPath.toLowerCase(); + String[] split = fullPath.split("\\."); + int i = 0; + boolean protocolEnd = false; + while (i < split.length) { + String s = split[i]; + boolean isL7Protocol = false; + if (l7ProtocolMap.containsKey(s)) { + isL7Protocol = true; + if (!protocolEnd) { + l7Protocol = l7ProtocolMap.get(s); + } + } + if (!isL7Protocol && !"UNCATEGORIZED".equals(l7Protocol)) { + protocolEnd = true; + } + if (i >= appIndex && !isL7Protocol) { + app = s; + } + i++; + } + } + event.getExtractedFields().put(l7ProtocolFieldName, l7Protocol); + event.getExtractedFields().put(appFieldName, app); + } catch (Exception e) { + logger.error("extract l7 protocol and app fail: " + e.getMessage(), e); + } + return event; + } + + @Override + public String functionName() { + return "CN_L7_PROTOCOL_AND_APP_EXTRACT"; + } + + @Override + public void close() { + + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/LinkDirectionLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/LinkDirectionLookup.java new file mode 100644 index 0000000..883ae01 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/LinkDirectionLookup.java @@ -0,0 +1,44 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.LinkDirectionKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/17 14:56 + */ +public class LinkDirectionLookup extends AbstractKnowledgeUDF { + + private LinkDirectionKnowledgeBaseHandler knowledgeBaseHandler; + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + event.getExtractedFields().put( + outputFieldName, + knowledgeBaseHandler.lookup( + event.getExtractedFields().get(lookupFieldName).toString() + ) + ); + } + return event; + } + + @Override + public String functionName() { + return "CN_LINK_DIRECTION_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + knowledgeBaseHandler = LinkDirectionKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(knowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/UserDefineTagLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/UserDefineTagLookup.java new file mode 100644 index 0000000..c10917d --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/UserDefineTagLookup.java @@ -0,0 +1,112 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.*; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; +import org.apache.flink.api.common.functions.RuntimeContext; + +import java.util.ArrayList; +import java.util.List; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 15:02 + */ +public class UserDefineTagLookup extends AbstractKnowledgeWithRuleUDF { + + private String option; + + private IpTagUserDefineKnowledgeBaseHandler ipKnowledgeBaseHandler; + private DomainTagUserDefineKnowledgeBaseHandler domainKnowledgeBaseHandler; + private AppTagUserDefineKnowledgeBaseHandler appKnowledgeBaseHandler; + private RuleKnowledgeBaseHandler ruleKnowledgeBaseHandler; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + option = udfContext.getParameters().get("option").toString(); + super.open(runtimeContext, udfContext); + } + + @Override + @SuppressWarnings("unchecked") + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + String lookupValue = event.getExtractedFields().get(lookupFieldName).toString(); + List<String> tags = new ArrayList<>(); + RuleMetadata ruleMetadata = new RuleMetadata(); + switch (option) { + case "IP_TO_TAG": + List<AbstractMultipleKnowledgeBaseHandler.Node> ipNodes = ipKnowledgeBaseHandler.lookup(lookupValue); + ipNodes.forEach(node -> { + tags.add(node.getTag()); + List<RuleKnowledgeBaseHandler.Rule> rules = ruleKnowledgeBaseHandler.lookupByKbId(node.getKbId()); + if (rules != null) { + rules.forEach(rule -> ruleMetadata.addRule(rule.getRuleId(), IocType.IP.getType())); + } + }); + break; + case "DOMAIN_TO_TAG": + List<AbstractMultipleKnowledgeBaseHandler.Node> domainNodes = domainKnowledgeBaseHandler.lookup(lookupValue); + domainNodes.forEach(node -> { + tags.add(node.getTag()); + List<RuleKnowledgeBaseHandler.Rule> rules = ruleKnowledgeBaseHandler.lookupByKbId(node.getKbId()); + if (rules != null) { + rules.forEach(rule -> ruleMetadata.addRule(rule.getRuleId(), IocType.DOMAIN.getType())); + } + }); + break; + case "APP_TO_TAG": + List<AbstractMultipleKnowledgeBaseHandler.Node> appNodes = appKnowledgeBaseHandler.lookup(lookupValue); + appNodes.forEach(node -> { + tags.add(node.getTag()); + List<RuleKnowledgeBaseHandler.Rule> rules = ruleKnowledgeBaseHandler.lookupByKbId(node.getKbId()); + if (rules != null) { + rules.forEach(rule -> ruleMetadata.addRule(rule.getRuleId(), IocType.APP.getType())); + } + }); + break; + default: + break; + } + event.getExtractedFields().put(outputFieldName, tags); + (((ArrayList<Long>) (event.getExtractedFields().computeIfAbsent(internalRuleIdListFieldName, k -> new ArrayList<Long>())))).addAll(ruleMetadata.getRuleIds()); + (((ArrayList<String>) (event.getExtractedFields().computeIfAbsent(internalIocTypeListFieldName, k -> new ArrayList<String>())))).addAll(ruleMetadata.getIocTypes()); + } + return event; + } + + @Override + public String functionName() { + return "CN_USER_DEFINE_TAG_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + ruleKnowledgeBaseHandler = RuleKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(ruleKnowledgeBaseHandler, ruleConfigs.get(0)); + + switch (option) { + case "IP_TO_TAG": + ipKnowledgeBaseHandler = IpTagUserDefineKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(ipKnowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + break; + case "DOMAIN_TO_TAG": + domainKnowledgeBaseHandler = DomainTagUserDefineKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(domainKnowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + break; + case "APP_TO_TAG": + appKnowledgeBaseHandler = AppTagUserDefineKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(appKnowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + break; + default: + break; + } + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/VpnLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/VpnLookup.java new file mode 100644 index 0000000..b084a03 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/VpnLookup.java @@ -0,0 +1,80 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.DomainVpnKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.IpVpnKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 10:40 + */ +public class VpnLookup extends AbstractKnowledgeUDF { + + private static final Logger logger = LoggerFactory.getLogger(VpnLookup.class); + + private String option; + + private IpVpnKnowledgeBaseHandler ipVpnKnowledgeBaseHandler; + + private DomainVpnKnowledgeBaseHandler domainVpnKnowledgeBaseHandler; + + @Override + public void open(RuntimeContext runtimeContext, UDFContext udfContext) { + option = udfContext.getParameters().get("option").toString(); + super.open(runtimeContext, udfContext); + } + + @Override + public Event evaluate(Event event) { + if (event.getExtractedFields().get(lookupFieldName) != null && event.getExtractedFields().get(lookupFieldName).toString().length() != 0) { + String lookup = event.getExtractedFields().get(lookupFieldName).toString(); + String result = null; + switch (option) { + case "IP_TO_VPN": + result = ipVpnKnowledgeBaseHandler.lookup(lookup); + break; + case "DOMAIN_TO_VPN": + result = domainVpnKnowledgeBaseHandler.lookup(lookup); + break; + default: + logger.error("unknown option: " + option); + break; + } + event.getExtractedFields().put(outputFieldName, result); + } + return event; + } + + @Override + public String functionName() { + return "CN_VPN_LOOKUP"; + } + + @Override + public void close() { + + } + + @Override + protected void registerKnowledges() { + switch (option) { + case "IP_TO_VPN": + ipVpnKnowledgeBaseHandler = IpVpnKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(ipVpnKnowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + break; + case "DOMAIN_TO_VPN": + domainVpnKnowledgeBaseHandler = DomainVpnKnowledgeBaseHandler.getInstance(); + KnowledgeBaseUpdateJob.registerKnowledgeBase(domainVpnKnowledgeBaseHandler, knowledgeBaseConfigs.get(0)); + break; + default: + break; + } + + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AbstractMultipleKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AbstractMultipleKnowledgeBaseHandler.java new file mode 100644 index 0000000..1d308a8 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AbstractMultipleKnowledgeBaseHandler.java @@ -0,0 +1,114 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.alibaba.fastjson2.JSON; +import com.geedgenetworks.common.config.KnowledgeBaseConfig; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.AbstractKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.cn.common.KnowledgeMetadata; +import lombok.AllArgsConstructor; +import lombok.Data; +import org.apache.http.HttpEntity; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.apache.http.util.EntityUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 15:34 + */ +public abstract class AbstractMultipleKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(AbstractMultipleKnowledgeBaseHandler.class); + + protected KnowledgeBaseConfig knowledgeBaseConfig; + + protected Map<String, KnowledgeMetadata> knowledgeMetedataCacheMap; + private static final CloseableHttpClient HTTP_CLIENT = HttpClients.createMinimal(); + + @Override + public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) { + this.knowledgeBaseConfig = knowledgeBaseConfig; + this.knowledgeMetedataCacheMap = getMetadata(knowledgeBaseConfig.getFsPath()); + return buildKnowledgeBase(); + } + + @Override + public void updateKnowledgeBase() { + if (ifNeedUpdate()) { + buildKnowledgeBase(); + } + } + + protected abstract Boolean buildKnowledgeBase(); + + protected Boolean ifNeedUpdate() { + Map<String, KnowledgeMetadata> knowledgeMetedataMap = getMetadata(knowledgeBaseConfig.getFsPath()); + if (knowledgeMetedataMap.size() != knowledgeMetedataCacheMap.size()) { + return true; + } + for (String id : knowledgeMetedataCacheMap.keySet()) { + if (knowledgeMetedataMap.containsKey(id) && knowledgeMetedataCacheMap.get(id).getSha256().equals(knowledgeMetedataMap.get(id).getSha256())) { + continue; + } + return true; + } + return false; + } + + public static Map<String, KnowledgeMetadata> getMetadata(String url) { + final HttpGet httpGet = new HttpGet(url); + httpGet.addHeader("Accept", "application/json"); + try { + CloseableHttpResponse response = HTTP_CLIENT.execute(httpGet); + HttpEntity entity = response.getEntity(); + if (entity != null) { + String content = EntityUtils.toString(entity, "UTF-8"); + KnowledgeResponse knowledgeResponse = JSON.parseObject(content, KnowledgeResponse.class); + List<KnowledgeMetadata> knowledgeMetedataList = JSON.parseArray(knowledgeResponse.data, KnowledgeMetadata.class); + return knowledgeMetedataList.stream() + .filter(metadata -> "latest".equals(metadata.getVersion()) && metadata.getIsValid() == 1 && metadata.getSha256() != null) + .collect(Collectors.toMap(KnowledgeMetadata::getId, Function.identity(), (existing, replacement) -> existing, HashMap::new)); + } + } catch (IOException e) { + logger.error("fetch knowledge metadata error", e); + } + return new HashMap<>(); + } + + @Data + private static final class KnowledgeResponse { + private int status; + private boolean success; + private String message; + private String data; + } + + /** + * 和知识库同步服务的转换规则保持一致 + * + * @param id + * @return + */ + protected Long convertId(String id) { + return Long.parseLong(id); + } + + @Data + @AllArgsConstructor + public static final class Node { + private String tag; + private Long kbId; + + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AbstractSingleKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AbstractSingleKnowledgeBaseHandler.java new file mode 100644 index 0000000..98f1ce9 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AbstractSingleKnowledgeBaseHandler.java @@ -0,0 +1,95 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.alibaba.fastjson2.JSON; +import com.geedgenetworks.common.config.KnowledgeBaseConfig; +import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.AbstractKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.cn.common.KnowledgeMetadata; +import lombok.Data; +import org.apache.http.HttpEntity; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.apache.http.util.EntityUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/18 17:42 + */ +public abstract class AbstractSingleKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(AbstractSingleKnowledgeBaseHandler.class); + + protected KnowledgeBaseConfig knowledgeBaseConfig; + + protected KnowLedgeBaseFileMeta knowledgeMetedataCache; + private static final CloseableHttpClient HTTP_CLIENT = HttpClients.createMinimal(); + + @Override + public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) { + this.knowledgeBaseConfig = knowledgeBaseConfig; + this.knowledgeMetedataCache = getMetadata(knowledgeBaseConfig.getFsType(), knowledgeBaseConfig.getFsPath(), knowledgeBaseConfig.getFiles().get(0)); + return buildKnowledgeBase(); + } + + @Override + public void updateKnowledgeBase() { + if (ifNeedUpdate()) { + buildKnowledgeBase(); + } + } + + protected abstract Boolean buildKnowledgeBase(); + + protected Boolean ifNeedUpdate() { + KnowLedgeBaseFileMeta knowledgeMetedata = getMetadata(knowledgeBaseConfig.getFsType(), knowledgeBaseConfig.getFsPath(), knowledgeBaseConfig.getFiles().get(0)); + if (knowledgeMetedata == null || knowledgeMetedata.getSha256() == null) { + return false; + } + if (knowledgeMetedataCache == null || knowledgeMetedataCache.getSha256() == null) { + knowledgeMetedataCache = knowledgeMetedata; + return true; + } + if (knowledgeMetedataCache.getSha256().equals(knowledgeMetedata.getSha256())) { + return false; + } else { + knowledgeMetedataCache = knowledgeMetedata; + return true; + } + } + + public List<KnowledgeMetadata> getMetadata(String url) { + final HttpGet httpGet = new HttpGet(url); + httpGet.addHeader("Accept", "application/json"); + try { + CloseableHttpResponse response = HTTP_CLIENT.execute(httpGet); + HttpEntity entity = response.getEntity(); + if (entity != null) { + String content = EntityUtils.toString(entity, "UTF-8"); + KnowledgeResponse knowledgeResponse = JSON.parseObject(content, KnowledgeResponse.class); + return JSON.parseArray(knowledgeResponse.data, KnowledgeMetadata.class).stream().filter(metadata -> "latest".equals(metadata.getVersion()) && metadata.getIsValid() == 1).collect(Collectors.toList()); + } + } catch (IOException e) { + logger.error("fetch knowledge metadata error", e); + } + return Collections.singletonList(null); + } + + @Data + private static final class KnowledgeResponse { + private int status; + private boolean success; + private String message; + private String data; + } + +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AppCategoryKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AppCategoryKnowledgeBaseHandler.java new file mode 100644 index 0000000..a9ae199 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AppCategoryKnowledgeBaseHandler.java @@ -0,0 +1,91 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import lombok.Data; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 14:08 + */ +public class AppCategoryKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(AppCategoryKnowledgeBaseHandler.class); + + private HashMap<String, AppCategory> appMap = new HashMap<>(); + + private AppCategoryKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final AppCategoryKnowledgeBaseHandler instance = new AppCategoryKnowledgeBaseHandler(); + } + + public static AppCategoryKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("app_name"); + needColumns.add("app_category"); + needColumns.add("app_subcategory"); + needColumns.add("app_company"); + needColumns.add("app_company_category"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<String, AppCategory> newAppMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String appName = line.get("app_name"); + String category = line.get("app_category"); + String subCategory = line.get("app_subcategory"); + String appCompany = line.get("app_company"); + String appCompanyCategory = line.get("app_company_category"); + + AppCategory appSketchLabel = new AppCategory(); + appSketchLabel.setAppName(appName); + appSketchLabel.setCategory(category); + appSketchLabel.setSubCategory(subCategory); + appSketchLabel.setAppCompany(appCompany); + appSketchLabel.setAppCompanyCategory(appCompanyCategory); + + newAppMap.put(appSketchLabel.getAppName(), appSketchLabel); + } catch (Exception lineException) { + logger.error("AppUtils line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + appMap = newAppMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public AppCategory lookup(String appName) { + return appMap.get(appName); + } + + @Data + public static final class AppCategory { + private String appName; + private String category; + private String subCategory; + private String appCompany; + private String appCompanyCategory; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AppTagUserDefineKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AppTagUserDefineKnowledgeBaseHandler.java new file mode 100644 index 0000000..f83596c --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/AppTagUserDefineKnowledgeBaseHandler.java @@ -0,0 +1,82 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.common.KnowledgeMetadata; +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 15:05 + */ +public class AppTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(AppTagUserDefineKnowledgeBaseHandler.class); + + private Map<String, List<Node>> appTagRules = new HashMap<>(); + + private AppTagUserDefineKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final AppTagUserDefineKnowledgeBaseHandler instance = new AppTagUserDefineKnowledgeBaseHandler(); + } + + public static AppTagUserDefineKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + Map<String, List<Node>> newAppTagRules = new HashMap<>(); + this.knowledgeMetedataCacheMap.forEach((key, value) -> { + buildSingleKnowledgeBase(newAppTagRules, convertId(key), value); + }); + appTagRules = newAppTagRules; + return true; + } + + private void buildSingleKnowledgeBase(Map<String, List<Node>> appTags, Long id, KnowledgeMetadata metadata) { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("tag_value"); + needColumns.add("app_name"); + byte[] content = downloadFile(metadata.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + String tagValue = line.get("tag_value"); + String appName = line.get("app_name"); + + if (appTags.containsKey(appName)) { + List<Node> nodes = appTags.get(appName); + nodes.add(new Node(tagValue, id)); + } else { + Node node = new Node(tagValue, id); + List<Node> nodes = new ArrayList<>(); + nodes.add(node); + appTags.put(appName, nodes); + } + } + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + } + } + + public List<Node> lookup(String app) { + if (appTagRules.containsKey(app)) { + return appTagRules.get(app); + } + return new ArrayList<>(); + } + +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DnsServerInfoKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DnsServerInfoKnowledgeBaseHandler.java new file mode 100644 index 0000000..5d12d8b --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DnsServerInfoKnowledgeBaseHandler.java @@ -0,0 +1,67 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/18 16:10 + */ +public class DnsServerInfoKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(DnsServerInfoKnowledgeBaseHandler.class); + + private HashMap<String, List<String>> dnsMap = new HashMap<>(); + + private DnsServerInfoKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final DnsServerInfoKnowledgeBaseHandler instance = new DnsServerInfoKnowledgeBaseHandler(); + } + + public static DnsServerInfoKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + public Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("ip_addr"); + needColumns.add("dns_server_role"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<String, List<String>> newDnsMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String ip = line.get("ip_addr"); + String dnsServerRole = line.get("dns_server_role"); + newDnsMap.computeIfAbsent(ip, k -> new ArrayList<String>()).add(dnsServerRole); + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + dnsMap = newDnsMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public List<String> lookup(String ip) { + return dnsMap.get(ip); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DomainTagUserDefineKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DomainTagUserDefineKnowledgeBaseHandler.java new file mode 100644 index 0000000..a9417e9 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DomainTagUserDefineKnowledgeBaseHandler.java @@ -0,0 +1,105 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.common.KnowledgeMetadata; +import com.geedgenetworks.core.utils.cn.common.Trie; +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.apache.commons.lang3.StringUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 15:05 + */ +public class DomainTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(DomainTagUserDefineKnowledgeBaseHandler.class); + + private Trie<Node> fqdnTagFuzzyMatchRules = new Trie<>(); + private Map<String, List<Node>> fqdnTagFullMatchRules = new HashMap<>(); + + private DomainTagUserDefineKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final DomainTagUserDefineKnowledgeBaseHandler instance = new DomainTagUserDefineKnowledgeBaseHandler(); + } + + public static DomainTagUserDefineKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + Trie<Node> newFqdnTagFuzzyMatchRules = new Trie<>(); + Map<String, List<Node>> newFqdnTagFullMatchRules = new HashMap<>(); + this.knowledgeMetedataCacheMap.forEach((key, value) -> { + buildSingleKnowledgeBase(newFqdnTagFuzzyMatchRules, newFqdnTagFullMatchRules, convertId(key), value); + }); + fqdnTagFuzzyMatchRules = newFqdnTagFuzzyMatchRules; + fqdnTagFullMatchRules = newFqdnTagFullMatchRules; + return true; + } + + private void buildSingleKnowledgeBase(Trie<Node> fqdnTagsFuzzy, Map<String, List<Node>> fqdnTagsFull, Long id, KnowledgeMetadata metadata) { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("tag_value"); + needColumns.add("domain"); + byte[] content = downloadFile(metadata.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + String tagValue = line.get("tag_value"); + String rule = line.get("domain"); + + if (rule == null || rule.length() == 0) { + continue; + } + if (!rule.startsWith("$") && !rule.startsWith("*")) { + continue; + } + String newRule = rule.substring(1); + if (rule.startsWith("*")) { + Node node = new Node(tagValue, id); + fqdnTagsFuzzy.put(StringUtils.reverse(newRule), node); + } else { + if (fqdnTagsFull.containsKey(newRule)) { + List<Node> nodes = fqdnTagsFull.get(newRule); + nodes.add(new Node(tagValue, id)); + } else { + Node node = new Node(tagValue, id); + List<Node> nodes = new ArrayList<>(); + nodes.add(node); + fqdnTagsFull.put(newRule, nodes); + } + } + } + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + } + } + + public List<Node> lookup(String domain) { + List<Node> nodes = new ArrayList<>(); + if (domain == null || domain.length() == 0) { + return nodes; + } + List<Node> fuzzyNodes = fqdnTagFuzzyMatchRules.get(StringUtils.reverse(domain)); + nodes.addAll(fuzzyNodes); + if (fqdnTagFullMatchRules.containsKey(domain)) { + List<Node> fullNodes = fqdnTagFullMatchRules.get(domain); + nodes.addAll(fullNodes); + } + return nodes; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DomainVpnKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DomainVpnKnowledgeBaseHandler.java new file mode 100644 index 0000000..5789283 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/DomainVpnKnowledgeBaseHandler.java @@ -0,0 +1,80 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 16:40 + */ +public class DomainVpnKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(DomainVpnKnowledgeBaseHandler.class); + + private Map<String, String> map = new HashMap<>(); + + private DomainVpnKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final DomainVpnKnowledgeBaseHandler instance = new DomainVpnKnowledgeBaseHandler(); + } + + public static DomainVpnKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("domain"); + needColumns.add("vpn_service_name"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + Map<String, String> newMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String domain = line.get("domain"); + String vpnServiceName = line.get("vpn_service_name"); + + newMap.put(domain, vpnServiceName); + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + map = newMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookup(String domain) { + if (domain == null || domain.length() == 0) { + return null; + } + if (map.containsKey(domain)) { + return map.get(domain); + } else { + int index = domain.indexOf(".") + 1; + if (index > 0) { + return lookup(domain.substring(index)); + } else { + return null; + } + } + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnCategoryKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnCategoryKnowledgeBaseHandler.java new file mode 100644 index 0000000..89d05ca --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnCategoryKnowledgeBaseHandler.java @@ -0,0 +1,99 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import lombok.Data; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 10:09 + */ +public class FqdnCategoryKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(FqdnCategoryKnowledgeBaseHandler.class); + + private HashMap<String, FqdnCategory> categoryMap = new HashMap<>(); + + private FqdnCategoryKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final FqdnCategoryKnowledgeBaseHandler instance = new FqdnCategoryKnowledgeBaseHandler(); + } + + public static FqdnCategoryKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("fqdn"); + needColumns.add("reputation_level"); + needColumns.add("category_name"); + needColumns.add("category_group"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<String, FqdnCategory> newCategoryMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String fqdn = line.get("fqdn"); + String reputationLevel = line.get("reputation_level"); + String categoryName = line.get("category_name"); + String categoryGroup = line.get("category_group"); + + FqdnCategory catalog = new FqdnCategory(); + catalog.setFqdn(fqdn); + catalog.setCategoryName(categoryName); + catalog.setCategoryGroup(categoryGroup); + catalog.setReputationLevel(reputationLevel); + + newCategoryMap.put(fqdn, catalog); + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " fqdn file line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + categoryMap = newCategoryMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public FqdnCategory lookup(String domain) { + if (domain == null || domain.length() == 0) { + return null; + } + if (categoryMap.containsKey(domain)) { + return categoryMap.get(domain); + } else { + int index = domain.indexOf(".") + 1; + if (index > 0) { + return lookup(domain.substring(index)); + } else { + return null; + } + } + } + + @Data + public static final class FqdnCategory { + private String fqdn; + private String categoryName; + private String categoryGroup; + private String reputationLevel; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnIcpKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnIcpKnowledgeBaseHandler.java new file mode 100644 index 0000000..685aaf4 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnIcpKnowledgeBaseHandler.java @@ -0,0 +1,79 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 10:28 + */ +public class FqdnIcpKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + private static final Logger logger = LoggerFactory.getLogger(FqdnIcpKnowledgeBaseHandler.class); + + private HashMap<String, String> icpMap = new HashMap<>(); + + private FqdnIcpKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final FqdnIcpKnowledgeBaseHandler instance = new FqdnIcpKnowledgeBaseHandler(); + } + + public static FqdnIcpKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("fqdn"); + needColumns.add("icp_company_name"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<String, String> newIcpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String fqdn = line.get("fqdn"); + String icpCompanyName = line.get("icp_company_name"); + + newIcpMap.put(fqdn, icpCompanyName); + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " icp file line: " + line.toString() + " parse error:" + lineException.getMessage(), lineException); + } + } + icpMap = newIcpMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookup(String domain) { + if (domain == null || domain.length() == 0) { + return null; + } + if (icpMap.containsKey(domain)) { + return icpMap.get(domain); + } else { + int index = domain.indexOf(".") + 1; + if (index > 0) { + return lookup(domain.substring(index)); + } else { + return null; + } + } + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnWhoisKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnWhoisKnowledgeBaseHandler.java new file mode 100644 index 0000000..f2ca3ff --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/FqdnWhoisKnowledgeBaseHandler.java @@ -0,0 +1,80 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 11:50 + */ +public class FqdnWhoisKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(FqdnWhoisKnowledgeBaseHandler.class); + + private HashMap<String, String> whoIsMap = new HashMap<>(); + + private FqdnWhoisKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final FqdnWhoisKnowledgeBaseHandler instance = new FqdnWhoisKnowledgeBaseHandler(); + } + + public static FqdnWhoisKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("fqdn"); + needColumns.add("whois_registrant_org"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<String, String> newWhoIsMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String fqdn = line.get("fqdn"); + String whoisRegistrantOrg = line.get("whois_registrant_org"); + + newWhoIsMap.put(fqdn, whoisRegistrantOrg); + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + whoIsMap = newWhoIsMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookup(String domain) { + if (domain == null || domain.length() == 0) { + return null; + } + if (whoIsMap.containsKey(domain)) { + return whoIsMap.get(domain); + } else { + int index = domain.indexOf(".") + 1; + if (index > 0) { + return lookup(domain.substring(index)); + } else { + return null; + } + } + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IdcRenterKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IdcRenterKnowledgeBaseHandler.java new file mode 100644 index 0000000..f180ce0 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IdcRenterKnowledgeBaseHandler.java @@ -0,0 +1,106 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.common.IPAddress; +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import inet.ipaddr.IPAddressString; +import org.apache.flink.shaded.guava18.com.google.common.collect.Range; +import org.apache.flink.shaded.guava18.com.google.common.collect.TreeRangeMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/18 11:22 + */ +public class IdcRenterKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(IdcRenterKnowledgeBaseHandler.class); + + private static TreeRangeMap<IPAddress, String> treeRangeMap = TreeRangeMap.create(); + + private IdcRenterKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final IdcRenterKnowledgeBaseHandler instance = new IdcRenterKnowledgeBaseHandler(); + } + + public static IdcRenterKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + public Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("addr_format"); + needColumns.add("ip1"); + needColumns.add("ip2"); + needColumns.add("idc_renter"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + TreeRangeMap<IPAddress, String> newTreeRangeMap = TreeRangeMap.create(); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String addrFormat = line.get("addr_format"); + String ip1 = line.get("ip1"); + String ip2 = line.get("ip2"); + String idcRenter = line.get("idc_renter"); + + if ("Single".equals(addrFormat)) { + IPAddress ipAddress = new IPAddress(ip1); + if (ipAddress.getIpAddress() == null) { + continue; + } + newTreeRangeMap.put(Range.closed(ipAddress, ipAddress), idcRenter); + } + if ("Range".equals(addrFormat)) { + IPAddress startIpAddress = new IPAddress(ip1); + IPAddress endIpAddress = new IPAddress(ip2); + if (startIpAddress.getIpAddress() == null || endIpAddress.getIpAddress() == null) { + continue; + } + newTreeRangeMap.put(Range.closed(startIpAddress, endIpAddress), idcRenter); + } + if ("CIDR".equals(addrFormat)) { + inet.ipaddr.IPAddress cidrIpAddress = new IPAddressString(ip1 + "/" + ip2).getAddress(); + if (cidrIpAddress == null) { + continue; + } + inet.ipaddr.IPAddress startIpAddress = cidrIpAddress.getLower(); + inet.ipaddr.IPAddress endIpAddress = cidrIpAddress.getUpper(); + newTreeRangeMap.put(Range.closed(new IPAddress(startIpAddress), new IPAddress(endIpAddress)), idcRenter); + } + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + treeRangeMap = newTreeRangeMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookup(String ip) { + try { + IPAddress address = new IPAddress(ip); + if (address.getIpAddress() != null) { + return treeRangeMap.get(address); + } + } catch (Exception e) { + logger.warn("invalid ip: " + ip); + } + return null; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/InternalIpKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/InternalIpKnowledgeBaseHandler.java new file mode 100644 index 0000000..d8cc44e --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/InternalIpKnowledgeBaseHandler.java @@ -0,0 +1,110 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.common.IPAddress; +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import inet.ipaddr.IPAddressString; +import org.apache.flink.shaded.guava18.com.google.common.collect.Range; +import org.apache.flink.shaded.guava18.com.google.common.collect.TreeRangeSet; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 17:46 + */ +public class InternalIpKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(InternalIpKnowledgeBaseHandler.class); + + private static TreeRangeSet<IPAddress> treeRangeSet = TreeRangeSet.create(); + + private InternalIpKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final InternalIpKnowledgeBaseHandler instance = new InternalIpKnowledgeBaseHandler(); + } + + public static InternalIpKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("addr_format"); + needColumns.add("ip1"); + needColumns.add("ip2"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + TreeRangeSet<IPAddress> newTreeRangeSet = TreeRangeSet.create(); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String addrFormat = line.get("addr_format"); + String ip1 = line.get("ip1"); + String ip2 = line.get("ip2"); + + if ("Single".equals(addrFormat)) { + IPAddress ipAddress = new IPAddress(ip1); + if (ipAddress.getIpAddress() == null) { + continue; + } + newTreeRangeSet.add(Range.closed(ipAddress, ipAddress)); + } + if ("Range".equals(addrFormat)) { + IPAddress startIpAddress = new IPAddress(ip1); + IPAddress endIpAddress = new IPAddress(ip2); + if (startIpAddress.getIpAddress() == null || endIpAddress.getIpAddress() == null) { + continue; + } + newTreeRangeSet.add(Range.closed(startIpAddress, endIpAddress)); + } + if ("CIDR".equals(addrFormat)) { + inet.ipaddr.IPAddress cidrIpAddress = new IPAddressString(ip1 + "/" + ip2).getAddress(); + if (cidrIpAddress == null) { + continue; + } + inet.ipaddr.IPAddress startIpAddress = cidrIpAddress.getLower(); + inet.ipaddr.IPAddress endIpAddress = cidrIpAddress.getUpper(); + newTreeRangeSet.add(Range.closed(new IPAddress(startIpAddress), new IPAddress(endIpAddress))); + } + } catch (Exception lineException) { + logger.error("IpUtils line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + treeRangeSet = newTreeRangeSet; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public boolean isInternal(String ip) { + try { + IPAddress ipAddress = new IPAddress(ip); + if (ipAddress.getIpAddress().isIPv4() && ipAddress.getIpAddress().toIPv4().isPrivate()) { + return true; + } + if (ipAddress.getIpAddress().isIPv6() && ipAddress.getIpAddress().toIPv6().isLocal()) { + return true; + } + if (treeRangeSet.contains(ipAddress)) { + return true; + } + } catch (Exception e) { + logger.error("internal ip "); + } + return false; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IocDarkwebKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IocDarkwebKnowledgeBaseHandler.java new file mode 100644 index 0000000..5898328 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IocDarkwebKnowledgeBaseHandler.java @@ -0,0 +1,93 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 10:53 + */ +public class IocDarkwebKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(IocDarkwebKnowledgeBaseHandler.class); + + private HashMap<String, String> ipMap = new HashMap<>(); + private HashMap<String, String> domainMap = new HashMap<>(); + + private IocDarkwebKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final IocDarkwebKnowledgeBaseHandler instance = new IocDarkwebKnowledgeBaseHandler(); + } + + public static IocDarkwebKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("ioc_type"); + needColumns.add("ioc_value"); + needColumns.add("node_type"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<String, String> newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HashMap<String, String> newDomainMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String type = line.get("ioc_type"); + String value = line.get("ioc_value"); + String nodeType = line.get("node_type"); + + if ("ip".equals(type)) { + newIpMap.put(value, nodeType); + } else if ("domain".equals(type)) { + newDomainMap.put(value, nodeType); + } + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + ipMap = newIpMap; + domainMap = newDomainMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookupByIp(String ip) { + return ipMap.get(ip); + } + + public String lookupByDomain(String domain) { + if (domain == null) { + return null; + } + if (domainMap.containsKey(domain)) { + return domainMap.get(domain); + } else { + int index = domain.indexOf(".") + 1; + if (index > 0) { + return lookupByDomain(domain.substring(index)); + } else { + return null; + } + } + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IocMalwareKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IocMalwareKnowledgeBaseHandler.java new file mode 100644 index 0000000..428ba1e --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IocMalwareKnowledgeBaseHandler.java @@ -0,0 +1,107 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.common.Trie; +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 14:46 + */ +public class IocMalwareKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(IocMalwareKnowledgeBaseHandler.class); + + private static HashMap<String, String> ipMap = new HashMap<>(); + private static HashMap<String, String> domainMap = new HashMap<>(); + private static Trie<String> urlTrie = new Trie<>(); + + private IocMalwareKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final IocMalwareKnowledgeBaseHandler instance = new IocMalwareKnowledgeBaseHandler(); + } + + public static IocMalwareKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("ioc_type"); + needColumns.add("ioc_value"); + needColumns.add("malware_name"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<String, String> newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HashMap<String, String> newDomainMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + Trie<String> newUrlTrie = new Trie<>(); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String type = line.get("ioc_type"); + String value = line.get("ioc_value"); + String malwareName = line.get("malware_name"); + + if ("ip".equals(type)) { + newIpMap.put(value, malwareName); + } else if ("domain".equals(type)) { + newDomainMap.put(value, malwareName); + } else if ("url".equals(type)) { + newUrlTrie.put(value, malwareName); + } + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + ipMap = newIpMap; + domainMap = newDomainMap; + urlTrie = newUrlTrie; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookupByIp(String ip) { + return ipMap.get(ip); + } + + public String lookupByDomain(String domain) { + if (domain == null) { + return null; + } + if (domainMap.containsKey(domain)) { + return domainMap.get(domain); + } else { + int index = domain.indexOf(".") + 1; + if (index > 0) { + return lookupByDomain(domain.substring(index)); + } else { + return null; + } + } + } + + public String lookupByUrl(String url) { + List<String> list = urlTrie.get(url); + if (list.size() == 0) { + return null; + } + return list.get(0); + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IpTagUserDefineKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IpTagUserDefineKnowledgeBaseHandler.java new file mode 100644 index 0000000..c6b4ac8 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IpTagUserDefineKnowledgeBaseHandler.java @@ -0,0 +1,127 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.common.KnowledgeMetadata; +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import inet.ipaddr.IPAddress; +import inet.ipaddr.IPAddressString; +import org.apache.flink.shaded.guava18.com.google.common.collect.Range; +import org.apache.flink.shaded.guava18.com.google.common.collect.TreeRangeMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 15:05 + */ +public class IpTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(IpTagUserDefineKnowledgeBaseHandler.class); + + private TreeRangeMap<IPAddress, List<Node>> ipTagRules = TreeRangeMap.create(); + + private IpTagUserDefineKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final IpTagUserDefineKnowledgeBaseHandler instance = new IpTagUserDefineKnowledgeBaseHandler(); + } + + public static IpTagUserDefineKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + TreeRangeMap<IPAddress, List<Node>> newIpTagRules = TreeRangeMap.create(); + this.knowledgeMetedataCacheMap.forEach((key, value) -> { + buildSingleKnowledgeBase(newIpTagRules, convertId(key), value); + }); + ipTagRules = newIpTagRules; + return true; + } + + private void buildSingleKnowledgeBase(TreeRangeMap<IPAddress, List<Node>> treeRangeMap, Long id, KnowledgeMetadata metadata) { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("tag_value"); + needColumns.add("addr_format"); + needColumns.add("ip1"); + needColumns.add("ip2"); + byte[] content = downloadFile(metadata.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String tagValue = line.get("tag_value"); + String addrFormat = line.get("addr_format"); + String ip1 = line.get("ip1"); + String ip2 = line.get("ip2"); + + IPAddress startIpAddress; + IPAddress endIpAddress; + if ("Single".equals(addrFormat)) { + IPAddress ipAddress = new IPAddressString(ip1).getAddress(); + if (ipAddress == null) { + continue; + } + startIpAddress = ipAddress; + endIpAddress = ipAddress; + } else if ("Range".equals(addrFormat)) { + IPAddress ipAddress1 = new IPAddressString(ip1).getAddress(); + IPAddress ipAddress2 = new IPAddressString(ip2).getAddress(); + if (ipAddress1 == null || ipAddress2 == null) { + continue; + } + startIpAddress = ipAddress1; + endIpAddress = ipAddress2; + } else if ("CIDR".equals(addrFormat)) { + IPAddress cidrIpAddress = new IPAddressString(ip1 + "/" + ip2).getAddress(); + if (cidrIpAddress == null) { + continue; + } + IPAddress ipAddressLower = cidrIpAddress.getLower(); + IPAddress ipAddressUpper = cidrIpAddress.getUpper(); + startIpAddress = ipAddressLower; + endIpAddress = ipAddressUpper; + } else { + logger.warn("unknown addrFormat: " + addrFormat); + continue; + } + + Map<Range<IPAddress>, List<Node>> rangeListMap = treeRangeMap.subRangeMap(Range.closed(startIpAddress, endIpAddress)).asMapOfRanges(); + TreeRangeMap<IPAddress, List<Node>> subRangeMap = TreeRangeMap.create(); + Node node = new Node(tagValue, id); + List<Node> nodes = new ArrayList<>(); + nodes.add(node); + subRangeMap.put(Range.closed(startIpAddress, endIpAddress), nodes); + rangeListMap.forEach((ipAddressRange, ipAddressNode) -> { + ipAddressNode.add(new Node(tagValue, id)); + subRangeMap.put(ipAddressRange, ipAddressNode); + }); + treeRangeMap.putAll(subRangeMap); + } catch (Exception lineException) { + logger.error("ip tag line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + } + } + + public List<Node> lookup(String ip) { + List<Node> nodes = null; + IPAddress address = new IPAddressString(ip).getAddress(); + if (address != null) { + nodes = ipTagRules.get(address); + } + return nodes == null ? new ArrayList<>() : nodes; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IpVpnKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IpVpnKnowledgeBaseHandler.java new file mode 100644 index 0000000..97eba37 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/IpVpnKnowledgeBaseHandler.java @@ -0,0 +1,106 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.common.IPAddress; +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import inet.ipaddr.IPAddressString; +import org.apache.flink.shaded.guava18.com.google.common.collect.Range; +import org.apache.flink.shaded.guava18.com.google.common.collect.TreeRangeMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/19 10:42 + */ +public class IpVpnKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(IpVpnKnowledgeBaseHandler.class); + + private TreeRangeMap<IPAddress, String> treeRangeMap = TreeRangeMap.create(); + + private IpVpnKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final IpVpnKnowledgeBaseHandler instance = new IpVpnKnowledgeBaseHandler(); + } + + public static IpVpnKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + protected Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("addr_format"); + needColumns.add("ip1"); + needColumns.add("ip2"); + needColumns.add("vpn_service_name"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + TreeRangeMap<IPAddress, String> newTreeRangeMap = TreeRangeMap.create(); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String addrFormat = line.get("addr_format"); + String ip1 = line.get("ip1"); + String ip2 = line.get("ip2"); + String vpnServiceName = line.get("vpn_service_name"); + + if ("Single".equals(addrFormat)) { + IPAddress ipAddress = new IPAddress(ip1); + if (ipAddress.getIpAddress() == null) { + continue; + } + newTreeRangeMap.put(Range.closed(ipAddress, ipAddress), vpnServiceName); + } + if ("Range".equals(addrFormat)) { + IPAddress startIpAddress = new IPAddress(ip1); + IPAddress endIpAddress = new IPAddress(ip2); + if (startIpAddress.getIpAddress() == null || endIpAddress.getIpAddress() == null) { + continue; + } + newTreeRangeMap.put(Range.closed(startIpAddress, endIpAddress), vpnServiceName); + } + if ("CIDR".equals(addrFormat)) { + inet.ipaddr.IPAddress cidrIpAddress = new IPAddressString(ip1 + "/" + ip2).getAddress(); + if (cidrIpAddress == null) { + continue; + } + inet.ipaddr.IPAddress startIpAddress = cidrIpAddress.getLower(); + inet.ipaddr.IPAddress endIpAddress = cidrIpAddress.getUpper(); + newTreeRangeMap.put(Range.closed(new IPAddress(startIpAddress), new IPAddress(endIpAddress)), vpnServiceName); + } + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + treeRangeMap = newTreeRangeMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookup(String ip) { + try { + IPAddress ipAddress = new IPAddress(ip); + if (ipAddress.getIpAddress() != null) { + return treeRangeMap.get(ipAddress); + } + } catch (Exception e) { + logger.error("ip vpn lookup faild: " + ip, e); + } + return null; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/LinkDirectionKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/LinkDirectionKnowledgeBaseHandler.java new file mode 100644 index 0000000..3ff9a6f --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/LinkDirectionKnowledgeBaseHandler.java @@ -0,0 +1,74 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.geedgenetworks.core.utils.cn.csv.HighCsvReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/17 15:06 + */ +public class LinkDirectionKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(LinkDirectionKnowledgeBaseHandler.class); + + private HashMap<Long, String> linkMap = new HashMap<>(); + + private LinkDirectionKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final LinkDirectionKnowledgeBaseHandler instance = new LinkDirectionKnowledgeBaseHandler(); + } + + public static LinkDirectionKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + public Boolean buildKnowledgeBase() { + try { + List<String> needColumns = new ArrayList<>(); + needColumns.add("peer_city"); + needColumns.add("link_id"); + byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1); + HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns); + HashMap<Long, String> newLinkMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); + HighCsvReader.CsvIterator iterator = highCsvReader.getIterator(); + while (iterator.hasNext()) { + Map<String, String> line = iterator.next(); + try { + String peerCity = line.get("peer_city"); + String linkIdStr = line.get("link_id"); + long linkId = Long.parseLong("".equals(linkIdStr) ? "0" : linkIdStr); + + newLinkMap.put(linkId, peerCity); + } catch (Exception lineException) { + logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error:" + lineException, lineException); + } + } + linkMap = newLinkMap; + } catch (Exception e) { + logger.error(this.getClass().getSimpleName() + " update error", e); + return false; + } + return true; + } + + public String lookup(String linkId) { + try { + long id = Long.parseLong(linkId); + return linkMap.get(id); + } catch (Exception e) { + logger.warn("invalid link id: " + linkId); + } + return null; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/RuleKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/RuleKnowledgeBaseHandler.java new file mode 100644 index 0000000..23ce14e --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/KnowlegdeBase/Handler/cn/RuleKnowledgeBaseHandler.java @@ -0,0 +1,181 @@ +package com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn; + +import com.alibaba.fastjson2.JSON; +import com.geedgenetworks.common.config.KnowledgeBaseConfig; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.AbstractKnowledgeBaseHandler; +import lombok.Data; +import org.apache.http.HttpEntity; +import org.apache.http.client.methods.CloseableHttpResponse; +import org.apache.http.client.methods.HttpGet; +import org.apache.http.impl.client.CloseableHttpClient; +import org.apache.http.impl.client.HttpClients; +import org.apache.http.util.EntityUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/22 16:35 + */ +public class RuleKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { + + private static final Logger logger = LoggerFactory.getLogger(RuleKnowledgeBaseHandler.class); + + protected KnowledgeBaseConfig knowledgeBaseConfig; + + private static final CloseableHttpClient HTTP_CLIENT = HttpClients.createMinimal(); + + private Long currentVersion; + + private Map<Long, Rule> ruleMap = new HashMap<>(); + + private Map<String, Rule> nameMap = new HashMap<>(); + + private Map<Long, List<Rule>> kbIdMap = new HashMap<>(); + + private RuleKnowledgeBaseHandler() { + } + + private static final class InstanceHolder { + private static final RuleKnowledgeBaseHandler instance = new RuleKnowledgeBaseHandler(); + } + + public static RuleKnowledgeBaseHandler getInstance() { + return InstanceHolder.instance; + } + + @Override + public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) { + this.knowledgeBaseConfig = knowledgeBaseConfig; + updateCache(knowledgeBaseConfig.getFsPath()); + return true; + } + + @Override + public void updateKnowledgeBase() { + updateCache(knowledgeBaseConfig.getFsPath()); + } + + public void updateCache(String address) { + String url = getUrl(address); + final HttpGet httpGet = new HttpGet(url); + httpGet.addHeader("Accept", "application/json"); + httpGet.addHeader("Cn-Authorization", knowledgeBaseConfig.getProperties().get("token")); + try { + CloseableHttpResponse response = HTTP_CLIENT.execute(httpGet); + HttpEntity entity = response.getEntity(); + if (entity != null) { + String content = EntityUtils.toString(entity, "UTF-8"); + RuleResponse ruleResponse = JSON.parseObject(content, RuleResponse.class); + processResponse(ruleResponse); + } + } catch (IOException e) { + logger.error("fetch rule metadata error", e); + } + } + + private void processResponse(RuleResponse ruleResponse) { + if (ruleResponse.data.version.equals(this.currentVersion)) { + return; + } + this.currentVersion = ruleResponse.data.version; + List<Rule> addList = ruleResponse.data.addList; + addList.addAll(ruleResponse.data.updateList); + List<Long> deleteIds = ruleResponse.data.deleteIds; + addList.stream().filter(rule -> rule.isBuiltIn == 1).forEach(rule -> { + nameMap.put(rule.name, rule); + ruleMap.put(rule.ruleId, rule); + }); + addList.stream().filter(rule -> rule.isBuiltIn != 1).forEach(rule -> { + kbIdMap.computeIfAbsent(rule.ruleConfigObj.getKnowledgeBase().getKnowledgeId(), k -> new ArrayList<Rule>()).add(rule); + ruleMap.put(rule.ruleId, rule); + }); + deleteIds.forEach(id -> { + Rule remove = ruleMap.remove(id); + if (remove.isBuiltIn == 1) { + nameMap.remove(remove.name); + } else { + Long knowledgeId = remove.ruleConfigObj.getKnowledgeBase().getKnowledgeId(); + if (kbIdMap.containsKey(knowledgeId)) { + List<Rule> rules = kbIdMap.get(knowledgeId); + rules.removeIf(next -> next.ruleId == remove.ruleId); + } + } + }); + } + + private String getUrl(String address) { + if (currentVersion == null) { + return address + "/v1/rule/detection"; + } + return address + "/v1/rule/detection/increase/" + currentVersion; + } + + public Rule lookupByName(String name) { + return nameMap.get(name); + } + + public List<Rule> lookupByKbId(Long kbId) { + return kbIdMap.get(kbId); + } + + @Data + private static final class RuleResponse { + private Integer code; + private String msg; + private DataValue data; + } + + @Data + private static final class DataValue { + private Long version; + private List<Rule> addList; + private List<Rule> updateList; + private List<Long> deleteIds; + } + + @Data + public static final class Rule { + private long ruleId; + private String ruleType; + private String name; + private String category; + private String eventType; + // Whether the rule is built-in (true) or custom (false). + private int isBuiltIn; + // Whether the current rule is enabled (true) or disabled (false). + private int status; + private RuleConfig ruleConfigObj; + private Trigger ruleTriggerObj; + } + + @Data + private static final class RuleConfig { + private String knowledgeId; + private KnowledgeBase knowledgeBase; + private String level; + private String dataSource; + } + + @Data + private static final class KnowledgeBase { + private Long knowledgeId; + private String source; + private String name; + private String category; + } + + @Data + private static final class Trigger { + private int atLeast; + private String interval; + private String resetInterval; + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/IPAddress.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/IPAddress.java new file mode 100644 index 0000000..e443e35 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/IPAddress.java @@ -0,0 +1,84 @@ +package com.geedgenetworks.core.utils.cn.common; + +import inet.ipaddr.AddressSection; +import inet.ipaddr.AddressStringException; +import inet.ipaddr.IPAddressString; +import inet.ipaddr.format.AddressDivisionSeries; +import inet.ipaddr.ipv4.IPv4AddressSection; +import inet.ipaddr.ipv6.IPv6AddressSection; + +public class IPAddress implements Comparable<IPAddress> { + + private inet.ipaddr.IPAddress ipAddress; + + public IPAddress(String ip) throws AddressStringException { + ipAddress = new IPAddressString(ip).getAddress(); + } + + public IPAddress(inet.ipaddr.IPAddress ipAddress) { + this.ipAddress = ipAddress; + } + + public inet.ipaddr.IPAddress getIpAddress() { + return ipAddress; + } + + public void setIpAddress(inet.ipaddr.IPAddress ipAddress) { + this.ipAddress = ipAddress; + } + + @Override + public int compareTo(IPAddress that) { + return compare(this.ipAddress.getSection(), that.ipAddress.getSection()); + } + + private int compare(AddressSection one, AddressSection two) { + if (one == two) { + return 0; + } + if (!one.getClass().equals(two.getClass())) { + int result = mapGrouping(one) - mapGrouping(two); + if (result != 0) { + return result; + } + } + if (one instanceof IPv6AddressSection) { + IPv6AddressSection o1 = (IPv6AddressSection) one; + IPv6AddressSection o2 = (IPv6AddressSection) two; + int result = o2.addressSegmentIndex - o1.addressSegmentIndex; + if (result != 0) { + return result; + } + } + + return compareParts(one, two); + } + + private int mapGrouping(AddressDivisionSeries series) { + if (series instanceof IPv6AddressSection) { + return 1; + } else if (series instanceof IPv4AddressSection) { + return -1; + } + return 0; + } + + private int compareParts(AddressSection one, AddressSection two) { + int segCount = one.getSegmentCount(); + for (int i = 0; i < segCount; i++) { + int oneLower = one.getSegment(i).getSegmentValue(); + int twoLower = two.getSegment(i).getSegmentValue(); + int result = compareValues(oneLower, twoLower); + if (result != 0) { + return result; + } + } + return 0; + } + + + protected int compareValues(int oneLower, int twoLower) { + return oneLower - twoLower; + } + +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/KnowledgeMetadata.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/KnowledgeMetadata.java new file mode 100644 index 0000000..6926852 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/KnowledgeMetadata.java @@ -0,0 +1,33 @@ +package com.geedgenetworks.core.utils.cn.common; + +import com.alibaba.fastjson2.annotation.JSONField; +import lombok.Data; +import lombok.ToString; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/18 17:35 + */ +@Data +@ToString +public class KnowledgeMetadata { + + @JSONField(name = "kb_id") + private String id; + + private String name; + + private String sha256; + + private String format; + + private String path; + + private String category; + + @JSONField(name = "is_valid") + private int isValid = 1; + + private String version; +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/Trie.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/Trie.java new file mode 100644 index 0000000..676815c --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/common/Trie.java @@ -0,0 +1,77 @@ +package com.geedgenetworks.core.utils.cn.common; + +import lombok.Data; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class Trie<T> { + + private final Node<T> root = new Node<>(); + + @Data + private static class Node<T> { + private char c; + private List<T> data = new ArrayList<>(); + private Map<Character, Node<T>> childNode = new HashMap<>(); + } + + public void put(String index, T data) { + List<T> list = new ArrayList<>(); + list.add(data); + put(root, index, list); + } + + public void putAll(String index, List<T> list) { + put(root, index, list); + } + + private void put(Node<T> node, String index, List<T> data) { + if (index == null || index.length() == 0) { + return; + } + boolean last = false; + if (index.length() == 1) { + last = true; + } + char c = index.charAt(0); + if (node.getChildNode().containsKey(c)) { + if (last) { + node.getChildNode().get(c).getData().addAll(data); + } else { + put(node.getChildNode().get(c), index.substring(1), data); + } + } else { + Node<T> newNode = new Node<T>(); + newNode.setC(c); + if (last) { + newNode.getData().addAll(data); + node.getChildNode().put(c, newNode); + } else { + node.getChildNode().put(c, newNode); + put(newNode, index.substring(1), data); + } + } + } + + public List<T> get(String str) { + List<T> result = new ArrayList<>(); + if (str == null || str.length() == 0) { + return result; + } + get(root, str, result); + return result; + } + + private void get(Node<T> node, String str, List<T> result) { + char c = str.charAt(0); + if (node.getChildNode().containsKey(c)) { + result.addAll(node.getChildNode().get(c).getData()); + if (str.length() > 1) { + get(node.getChildNode().get(c), str.substring(1), result); + } + } + } +} diff --git a/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/csv/HighCsvReader.java b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/csv/HighCsvReader.java new file mode 100644 index 0000000..7f4a9d6 --- /dev/null +++ b/groot-core/src/main/java/com/geedgenetworks/core/utils/cn/csv/HighCsvReader.java @@ -0,0 +1,123 @@ +package com.geedgenetworks.core.utils.cn.csv; + +import com.opencsv.CSVParser; +import com.opencsv.CSVReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.ByteArrayInputStream; +import java.io.InputStreamReader; +import java.io.Reader; +import java.util.*; + +public class HighCsvReader { + + private final Logger logger = LoggerFactory.getLogger(HighCsvReader.class); + + private CSVReader csvReader; + + private String[] header = null; + + private final Map<String, Integer> columnNames = new HashMap<>(); + + private final List<String[]> data = new ArrayList<>(); + + public HighCsvReader(Reader reader) { + init(reader, new ArrayList<>()); + } + + public HighCsvReader(Reader reader, List<String> columns) { + init(reader, columns); + } + + private void init(Reader reader, List<String> columns) { + try { + this.csvReader = new CSVReader(reader, CSVParser.DEFAULT_SEPARATOR, + CSVParser.DEFAULT_QUOTE_CHARACTER, CSVParser.DEFAULT_ESCAPE_CHARACTER); + String[] allHeader = csvReader.readNext(); + if (allHeader == null) { + return; + } + for (int i = 0; i < allHeader.length; i++) { + String columnName = allHeader[i]; + columnNames.put(columnName, i); + } + if (columns.size() == 0) { + columns.addAll(columnNames.keySet()); + } else { + columns.removeIf(next -> !columnNames.containsKey(next)); + } + header = new String[columns.size()]; + columns.toArray(header); + + for (String[] strings : csvReader) { + String[] lineData = new String[header.length]; + for (int i = 0; i < lineData.length; i++) { + lineData[i] = strings[columnNames.get(header[i])]; + } + data.add(lineData); + } + } catch (Exception e) { + logger.error("csv parse error, error exception: " + e.getMessage(), e); + } finally { + if (csvReader != null) { + try { + csvReader.close(); + } catch (Exception ignored) { + + } + } + } + } + + public CsvIterator getIterator() { + return new CsvIterator(header, data); + } + + public int getLineNumber() { + return data.size(); + } + + public Integer getFieldIndex(String fieldName) { + return columnNames.get(fieldName); + } + + private InputStreamReader getReader(byte[] bytes) { + return new InputStreamReader(new ByteArrayInputStream(bytes)); + } + + public class CsvIterator implements Iterator<Map<String, String>> { + + private String[] header; + + private List<String[]> data; + + private int total = 0; + + private int currentLineNumber = 0; + + public CsvIterator(String[] header, List<String[]> data) { + this.header = header; + this.data = data; + this.total = data.size(); + } + + @Override + public boolean hasNext() { + return currentLineNumber < total; + } + + @Override + public Map<String, String> next() { + HashMap<String, String> lineMap = new HashMap<String, String>(header.length); + String[] line = data.get(currentLineNumber); + for (int i = 0; i < line.length; i++) { + lineMap.put(header[i], line[i]); + } + currentLineNumber++; + return lineMap; + } + } + + +} diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/AnonymityLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/AnonymityLookupTest.java new file mode 100644 index 0000000..a52deb1 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/AnonymityLookupTest.java @@ -0,0 +1,75 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class AnonymityLookupTest { + + private static AnonymityLookup anonymityLookup; + + private static RuntimeContext runtimeContext; + + @BeforeEach + void setUp() { + runtimeContext = mockRuleRuntimeContext(); + + String content = "ioc_type,ioc_value,node_type\nip,92.230.220.168,tor\ndomain,v103.vipgorup.xyz,mtproxy"; + mockKnowledgeBaseHandler(content); + + anonymityLookup = new AnonymityLookup(); + } + + @Test + void evaluateIp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "IP_TO_NODE_TYPE"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("server_ip")); + udfContext.setOutput_fields(Collections.singletonList("server_node_type")); + anonymityLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("server_ip", "92.230.220.168"); + event.setExtractedFields(fields); + Event evaluate = anonymityLookup.evaluate(event); + assertEquals("tor", evaluate.getExtractedFields().get("server_node_type")); + } + + @Test + void evaluateDomain() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "DOMAIN_TO_NODE_TYPE"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + udfContext.setOutput_fields(Collections.singletonList("domain_node_type")); + anonymityLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "v103.vipgorup.xyz"); + event.setExtractedFields(fields); + Event evaluate = anonymityLookup.evaluate(event); + assertEquals("mtproxy", evaluate.getExtractedFields().get("domain_node_type")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/AppCategoryLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/AppCategoryLookupTest.java new file mode 100644 index 0000000..f74ce39 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/AppCategoryLookupTest.java @@ -0,0 +1,61 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class AppCategoryLookupTest { + + private static AppCategoryLookup appCategoryLookup; + + @BeforeAll + static void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + Map<String, String> fieldMapping = new HashMap<>(); + fieldMapping.put("CATEGORY", "app_category"); + fieldMapping.put("SUBCATEGORY", "app_subcategory"); + fieldMapping.put("COMPANY", "app_company"); + fieldMapping.put("COMPANY_CATEGORY", "app_company_category"); + parameters.put("field_mapping", fieldMapping); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("app")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "app_name,app_category,app_subcategory,app_company,app_company_category\nqq,collaboration,instant-messaging,Tencent,TencentCategory"; + mockKnowledgeBaseHandler(content); + + appCategoryLookup = new AppCategoryLookup(); + appCategoryLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluate() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("app", "qq"); + event.setExtractedFields(fields); + Event evaluate = appCategoryLookup.evaluate(event); + assertEquals("collaboration", evaluate.getExtractedFields().get("app_category")); + assertEquals("instant-messaging", evaluate.getExtractedFields().get("app_subcategory")); + assertEquals("Tencent", evaluate.getExtractedFields().get("app_company")); + assertEquals("TencentCategory", evaluate.getExtractedFields().get("app_company_category")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/DnsServerInfoLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/DnsServerInfoLookupTest.java new file mode 100644 index 0000000..7f526d5 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/DnsServerInfoLookupTest.java @@ -0,0 +1,54 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class DnsServerInfoLookupTest { + + private static DnsServerInfoLookup dnsServerInfoLookup; + + @BeforeAll + static void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("server_ip")); + udfContext.setOutput_fields(Collections.singletonList("server_dns_server")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "ip_addr,dns_server_role\n111.196.123.207,FWDNS"; + mockKnowledgeBaseHandler(content); + + dnsServerInfoLookup = new DnsServerInfoLookup(); + dnsServerInfoLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluate() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("server_ip", "111.196.123.207"); + event.setExtractedFields(fields); + Event evaluate = dnsServerInfoLookup.evaluate(event); + assertEquals(Arrays.asList("FWDNS"), evaluate.getExtractedFields().get("server_dns_server")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/FqdnCategoryLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/FqdnCategoryLookupTest.java new file mode 100644 index 0000000..db15642 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/FqdnCategoryLookupTest.java @@ -0,0 +1,59 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class FqdnCategoryLookupTest { + + private static FqdnCategoryLookup fqdnCategoryLookup; + + @BeforeAll + static void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + Map<String, String> fieldMapping = new HashMap<>(); + fieldMapping.put("NAME", "domain_category_name"); + fieldMapping.put("GROUP", "domain_category_group"); + fieldMapping.put("REPUTATION_LEVEL", "domain_reputation_level"); + parameters.put("field_mapping", fieldMapping); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "fqdn,category_name,category_group,reputation_level\ndouralquran.com,Parked Domains,Sensitive,Trustworthy"; + mockKnowledgeBaseHandler(content); + + fqdnCategoryLookup = new FqdnCategoryLookup(); + fqdnCategoryLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluate() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "www.douralquran.com"); + event.setExtractedFields(fields); + Event evaluate = fqdnCategoryLookup.evaluate(event); + assertEquals("Parked Domains", evaluate.getExtractedFields().get("domain_category_name")); + assertEquals("Sensitive", evaluate.getExtractedFields().get("domain_category_group")); + assertEquals("Trustworthy", evaluate.getExtractedFields().get("domain_reputation_level")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/FqdnWhoisLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/FqdnWhoisLookupTest.java new file mode 100644 index 0000000..42a98dc --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/FqdnWhoisLookupTest.java @@ -0,0 +1,53 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class FqdnWhoisLookupTest { + + private static FqdnWhoisLookup fqdnWhoisLookup; + + @BeforeAll + static void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + udfContext.setOutput_fields(Collections.singletonList("domain_whois_org")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "fqdn,whois_registrant_org\nschoolexpress.com,comped"; + mockKnowledgeBaseHandler(content); + + fqdnWhoisLookup = new FqdnWhoisLookup(); + fqdnWhoisLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluate() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "www.schoolexpress.com"); + event.setExtractedFields(fields); + Event evaluate = fqdnWhoisLookup.evaluate(event); + assertEquals("comped", evaluate.getExtractedFields().get("domain_whois_org")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IcpLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IcpLookupTest.java new file mode 100644 index 0000000..3158124 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IcpLookupTest.java @@ -0,0 +1,53 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class IcpLookupTest { + + private static IcpLookup icpLookup; + + @BeforeAll + static void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + udfContext.setOutput_fields(Collections.singletonList("domain_icp_company_name")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "fqdn,icp_company_name\ndreamtec-solutions.com.cn,上海骏品电子科技有限公司"; + mockKnowledgeBaseHandler(content); + + icpLookup = new IcpLookup(); + icpLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluate() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "www.dreamtec-solutions.com.cn"); + event.setExtractedFields(fields); + Event evaluate = icpLookup.evaluate(event); + assertEquals("上海骏品电子科技有限公司", evaluate.getExtractedFields().get("domain_icp_company_name")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IdcRenterLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IdcRenterLookupTest.java new file mode 100644 index 0000000..b15096b --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IdcRenterLookupTest.java @@ -0,0 +1,53 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class IdcRenterLookupTest { + + private static IdcRenterLookup idcRenterLookup; + + @BeforeAll + static void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("server_ip")); + udfContext.setOutput_fields(Collections.singletonList("server_idc_renter")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "addr_format,ip1,ip2,idc_renter\nCIDR,116.178.65.0,25,阿里"; + mockKnowledgeBaseHandler(content); + + idcRenterLookup = new IdcRenterLookup(); + idcRenterLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluate() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("server_ip", "116.178.65.100"); + event.setExtractedFields(fields); + Event evaluate = idcRenterLookup.evaluate(event); + assertEquals("阿里", evaluate.getExtractedFields().get("server_idc_renter")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IocLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IocLookupTest.java new file mode 100644 index 0000000..f9d3b25 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IocLookupTest.java @@ -0,0 +1,75 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class IocLookupTest { + + private static IocLookup iocLookup; + + private static RuntimeContext runtimeContext; + + @BeforeEach + void setUp() { + runtimeContext = mockRuleRuntimeContext(); + + String content = "ioc_type,ioc_value,malware_name\nip,2.50.159.104,QakBot\ndomain,72ioey.badomininter.cloud,Astaroth"; + mockKnowledgeBaseHandler(content); + + iocLookup = new IocLookup(); + } + + @Test + void evaluateIp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "IP_TO_MALWARE"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("server_ip")); + udfContext.setOutput_fields(Collections.singletonList("server_malware")); + iocLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("server_ip", "2.50.159.104"); + event.setExtractedFields(fields); + Event evaluate = iocLookup.evaluate(event); + assertEquals("QakBot", evaluate.getExtractedFields().get("server_malware")); + } + + @Test + void evaluateDomain() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "DOMAIN_TO_MALWARE"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + udfContext.setOutput_fields(Collections.singletonList("domain_malware")); + iocLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "test.72ioey.badomininter.cloud"); + event.setExtractedFields(fields); + Event evaluate = iocLookup.evaluate(event); + assertEquals("Astaroth", evaluate.getExtractedFields().get("domain_malware")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IpZoneLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IpZoneLookupTest.java new file mode 100644 index 0000000..abe5ba0 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IpZoneLookupTest.java @@ -0,0 +1,63 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class IpZoneLookupTest { + + private static IpZoneLookup ipZoneLookup; + + @BeforeEach + void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("client_ip")); + udfContext.setOutput_fields(Collections.singletonList("client_zone")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "addr_format,ip1,ip2\nCIDR,116.178.65.0,25"; + mockKnowledgeBaseHandler(content); + + ipZoneLookup = new IpZoneLookup(); + ipZoneLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluateInternal() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("client_ip", "116.178.65.100"); + event.setExtractedFields(fields); + Event evaluate = ipZoneLookup.evaluate(event); + assertEquals("internal", evaluate.getExtractedFields().get("client_zone")); + } + + @Test + void evaluateExternal() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("client_ip", "116.178.64.100"); + event.setExtractedFields(fields); + Event evaluate = ipZoneLookup.evaluate(event); + assertEquals("external", evaluate.getExtractedFields().get("client_zone")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LinkDirectionLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LinkDirectionLookupTest.java new file mode 100644 index 0000000..c0a06fe --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LinkDirectionLookupTest.java @@ -0,0 +1,53 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class LinkDirectionLookupTest { + + private static LinkDirectionLookup linkDirectionLookup; + + @BeforeAll + static void setUp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("in_link_id")); + udfContext.setOutput_fields(Collections.singletonList("in_link_direction")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "peer_city,link_id\n北京,1"; + mockKnowledgeBaseHandler(content); + + linkDirectionLookup = new LinkDirectionLookup(); + linkDirectionLookup.open(runtimeContext, udfContext); + } + + @Test + void evaluate() { + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("in_link_id", "1"); + event.setExtractedFields(fields); + Event evaluate = linkDirectionLookup.evaluate(event); + assertEquals("北京", evaluate.getExtractedFields().get("in_link_direction")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java new file mode 100644 index 0000000..6678e2e --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java @@ -0,0 +1,115 @@ +package com.geedgenetworks.core.udf.cn; + +import com.alibaba.fastjson2.JSON; +import com.geedgenetworks.common.Constants; +import com.geedgenetworks.common.config.CommonConfig; +import com.geedgenetworks.common.config.KnowledgeBaseConfig; +import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.AbstractKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.Handler.cn.AbstractMultipleKnowledgeBaseHandler; +import com.geedgenetworks.core.utils.KnowlegdeBase.KnowledgeBaseUpdateJob; +import com.geedgenetworks.core.utils.cn.common.KnowledgeMetadata; +import org.apache.flink.api.common.ExecutionConfig; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.apache.flink.configuration.Configuration; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.mockito.Mockito.mockStatic; + +/** + * @author gujinkai + * @version 1.0 + * @date 2024/1/26 16:45 + */ +public class LookupTestUtils { + + private static String fsPath = "testPath"; + private static String fsType = "testType"; + private static List<String> fsFiles = Arrays.asList("testFile"); + public static String kbName = "testKbName"; + private static String downloadPath = "testDownloadPath"; + + private static MockedStatic<AbstractKnowledgeBaseHandler> abstractKnowledgeBaseHandlerMockedStatic = mockStatic(AbstractKnowledgeBaseHandler.class); + + private static MockedStatic<AbstractMultipleKnowledgeBaseHandler> abstractMultipleKnowledgeBaseHandlerMockedStatic = mockStatic(AbstractMultipleKnowledgeBaseHandler.class); + + static RuntimeContext mockRuntimeContext() { + RuntimeContext runtimeContext = Mockito.mock(RuntimeContext.class); + ExecutionConfig executionConfig = Mockito.mock(ExecutionConfig.class); + Mockito.when(runtimeContext.getExecutionConfig()).thenReturn(executionConfig); + Configuration configuration = new Configuration(); + CommonConfig commonConfig = new CommonConfig(); + KnowledgeBaseConfig knowledgeBaseConfig = new KnowledgeBaseConfig(); + knowledgeBaseConfig.setFsPath(fsPath); + knowledgeBaseConfig.setFsType(fsType); + knowledgeBaseConfig.setFiles(fsFiles); + knowledgeBaseConfig.setName(kbName); + commonConfig.setKnowledgeBaseConfig(Arrays.asList(knowledgeBaseConfig)); + configuration.setString(Constants.SYSPROP_GROOTSTREAM_CONFIG, JSON.toJSONString(commonConfig)); + Mockito.when(executionConfig.getGlobalJobParameters()).thenReturn(configuration); + return runtimeContext; + } + + static void checkStaticMock() { + if (abstractKnowledgeBaseHandlerMockedStatic.isClosed()) { + abstractKnowledgeBaseHandlerMockedStatic = mockStatic(AbstractKnowledgeBaseHandler.class); + } + if (abstractMultipleKnowledgeBaseHandlerMockedStatic.isClosed()) { + abstractMultipleKnowledgeBaseHandlerMockedStatic = mockStatic(AbstractMultipleKnowledgeBaseHandler.class); + } + } + + static void mockKnowledgeBaseHandler(String downloadContent) { + checkStaticMock(); + KnowLedgeBaseFileMeta knowLedgeBaseFileMeta = new KnowLedgeBaseFileMeta(); + knowLedgeBaseFileMeta.setPath(downloadPath); + abstractKnowledgeBaseHandlerMockedStatic.when(() -> AbstractKnowledgeBaseHandler.getMetadata(fsType, fsPath, fsFiles.get(0))).thenReturn(knowLedgeBaseFileMeta); + abstractKnowledgeBaseHandlerMockedStatic.when(() -> AbstractKnowledgeBaseHandler.downloadFile(downloadPath, 1)).thenReturn(downloadContent.getBytes()); + } + + static void mockMultipleKnowledgeBaseHandler(String downloadContent) { + checkStaticMock(); + KnowledgeMetadata knowledgeMetadata = new KnowledgeMetadata(); + knowledgeMetadata.setPath(downloadPath); + Map<String, KnowledgeMetadata> knowledgeMetadataMap = new HashMap<>(); + knowledgeMetadataMap.put("1", knowledgeMetadata); + abstractMultipleKnowledgeBaseHandlerMockedStatic.when(() -> AbstractMultipleKnowledgeBaseHandler.getMetadata(fsPath)).thenReturn(knowledgeMetadataMap); + abstractKnowledgeBaseHandlerMockedStatic.when(() -> AbstractKnowledgeBaseHandler.downloadFile(downloadPath, 1)).thenReturn(downloadContent.getBytes()); + } + + static RuntimeContext mockRuleRuntimeContext() { + RuntimeContext runtimeContext = Mockito.mock(RuntimeContext.class); + ExecutionConfig executionConfig = Mockito.mock(ExecutionConfig.class); + Mockito.when(runtimeContext.getExecutionConfig()).thenReturn(executionConfig); + Configuration configuration = new Configuration(); + CommonConfig commonConfig = new CommonConfig(); + KnowledgeBaseConfig knowledgeBaseConfig = new KnowledgeBaseConfig(); + knowledgeBaseConfig.setFsPath(fsPath); + knowledgeBaseConfig.setFsType(fsType); + knowledgeBaseConfig.setFiles(fsFiles); + knowledgeBaseConfig.setName(kbName); + KnowledgeBaseConfig ruleKnowledgeBaseConfig = new KnowledgeBaseConfig(); + ruleKnowledgeBaseConfig.setFsPath("testRulePath"); + ruleKnowledgeBaseConfig.setFsType("testRuleType"); + Map<String, String> properties = new HashMap<>(); + properties.put("token", "testToken"); + ruleKnowledgeBaseConfig.setProperties(properties); + ruleKnowledgeBaseConfig.setName("cn_rule"); + commonConfig.setKnowledgeBaseConfig(Arrays.asList(knowledgeBaseConfig, ruleKnowledgeBaseConfig)); + configuration.setString(Constants.SYSPROP_GROOTSTREAM_CONFIG, JSON.toJSONString(commonConfig)); + Mockito.when(executionConfig.getGlobalJobParameters()).thenReturn(configuration); + return runtimeContext; + } + + static void clearState() { + KnowledgeBaseUpdateJob.removeKnowledgeBase(kbName); + abstractKnowledgeBaseHandlerMockedStatic.close(); + abstractMultipleKnowledgeBaseHandlerMockedStatic.close(); + } +} diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/TrieTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/TrieTest.java new file mode 100644 index 0000000..b54d13d --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/TrieTest.java @@ -0,0 +1,102 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.core.utils.cn.common.Trie; +import org.apache.commons.lang3.StringUtils; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; + + +public class TrieTest { + + @Test + public void TrieGetTest() { + Trie<String> trie = new Trie<>(); + trie.put(StringUtils.reverse("baidu.com"), "1"); + trie.put(StringUtils.reverse("test.baidu.com"), "2"); + trie.put(StringUtils.reverse("test1.baidu.com"), "3"); + trie.put(StringUtils.reverse("txj/r~/moc.elgoog.yxorpdeef//:ptth"), "4"); + + List<String> strings1 = trie.get(StringUtils.reverse("test.baidu.com")); + assertEquals(Arrays.asList("1", "2"), strings1); + + List<String> strings2 = trie.get(StringUtils.reverse("baidu.com")); + assertEquals(Arrays.asList("1"), strings2); + + List<String> strings3 = trie.get(StringUtils.reverse("test.test.baidu.com")); + assertEquals(Arrays.asList("1", "2"), strings3); + + List<String> strings4 = trie.get(StringUtils.reverse("test1.test.baidu.com")); + assertEquals(Arrays.asList("1", "2"), strings4); + + List<String> strings5 = trie.get(StringUtils.reverse("test1.test1.baidu.com")); + assertEquals(Arrays.asList("1", "3"), strings5); + + List<String> strings6 = trie.get(StringUtils.reverse("")); + assertEquals(Arrays.asList(), strings6); + + List<String> strings7 = trie.get(StringUtils.reverse(null)); + assertEquals(Arrays.asList(), strings7); + + List<String> strings8 = trie.get(StringUtils.reverse("txj/r~/moc.elgoog.yxorpdeef//:ptth")); + assertEquals(Arrays.asList("4"), strings8); + } + + @Test + public void testPutAndGet() { + Trie<String> trie = new Trie<>(); + + // Test putting a single data + trie.put("apple", "fruit"); + assertEquals(Arrays.asList("fruit"), trie.get("apple")); + + // Test putting multiple data + trie.putAll("banana", Arrays.asList("fruit", "yellow")); + assertEquals(Arrays.asList("fruit", "yellow"), trie.get("banana")); + + // Test putting data with the same prefix + trie.put("app", "application"); + trie.put("app", "appendix"); + assertEquals(Arrays.asList("application", "appendix"), trie.get("app")); + + // Test putting and getting with empty string + trie.put("", "empty"); + assertEquals(Arrays.asList(), trie.get("")); + } + + @Test + public void testPutAll() { + Trie<Integer> trie = new Trie<>(); + + // Test putting multiple data at once + trie.putAll("one", Arrays.asList(1, 11, 111)); + assertEquals(Arrays.asList(1, 11, 111), trie.get("one")); + + // Test putting an empty list + trie.putAll("empty", Arrays.asList()); + assertEquals(Arrays.asList(), trie.get("empty")); + } + + @Test + public void testGet() { + Trie<Double> trie = new Trie<>(); + + // Test getting data with non-existent prefix + assertEquals(Arrays.asList(), trie.get("nonexistent")); + + // Test getting data with partially matched prefix + trie.put("test", 3.14); + trie.put("testing", 2.71); + assertEquals(Arrays.asList(), trie.get("tes")); + + // Test getting data with fully matched prefix + assertEquals(Arrays.asList(3.14), trie.get("test")); + + // Test getting data with empty string + trie.put("empty", 0.0); + assertEquals(Arrays.asList(), trie.get("")); + } +} diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/UserDefineTagLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/UserDefineTagLookupTest.java new file mode 100644 index 0000000..9688199 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/UserDefineTagLookupTest.java @@ -0,0 +1,122 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class UserDefineTagLookupTest { + private static UserDefineTagLookup userDefineTagLookup; + + private static RuntimeContext runtimeContext; + + @BeforeEach + void setUp() { + runtimeContext = mockRuleRuntimeContext(); + + userDefineTagLookup = new UserDefineTagLookup(); + } + + @Test + void evaluateIp() { + String content = "addr_format,ip1,ip2,tag_value\nCIDR,103.3.136.128,25,4G"; + mockMultipleKnowledgeBaseHandler(content); + + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "IP_TO_TAG"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("client_ip")); + udfContext.setOutput_fields(Collections.singletonList("client_ip_tags")); + userDefineTagLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("client_ip", "103.3.136.130"); + event.setExtractedFields(fields); + Event evaluate = userDefineTagLookup.evaluate(event); + assertEquals(Arrays.asList("4G"), evaluate.getExtractedFields().get("client_ip_tags")); + } + + @Test + void evaluateDomainFuzzy() { + String content = "domain,tag_value\n*baidu.com,test_tag"; + mockMultipleKnowledgeBaseHandler(content); + + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "DOMAIN_TO_TAG"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + udfContext.setOutput_fields(Collections.singletonList("domain_tags")); + userDefineTagLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "www.baidu.com"); + event.setExtractedFields(fields); + Event evaluate = userDefineTagLookup.evaluate(event); + assertEquals(Arrays.asList("test_tag"), evaluate.getExtractedFields().get("domain_tags")); + } + + @Test + void evaluateDomainFull() { + String content = "domain,tag_value\n$baidu.com,test_tag"; + mockMultipleKnowledgeBaseHandler(content); + + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "DOMAIN_TO_TAG"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + udfContext.setOutput_fields(Collections.singletonList("domain_tags")); + userDefineTagLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "baidu.com"); + event.setExtractedFields(fields); + Event evaluate = userDefineTagLookup.evaluate(event); + assertEquals(Arrays.asList("test_tag"), evaluate.getExtractedFields().get("domain_tags")); + } + + @Test + void evaluateApp() { + String content = "app_name,tag_value\nbaidu,test_tag"; + mockMultipleKnowledgeBaseHandler(content); + + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "APP_TO_TAG"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("app")); + udfContext.setOutput_fields(Collections.singletonList("app_tags")); + userDefineTagLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("app", "baidu"); + event.setExtractedFields(fields); + Event evaluate = userDefineTagLookup.evaluate(event); + assertEquals(Arrays.asList("test_tag"), evaluate.getExtractedFields().get("app_tags")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/VpnLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/VpnLookupTest.java new file mode 100644 index 0000000..3dd0992 --- /dev/null +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/VpnLookupTest.java @@ -0,0 +1,80 @@ +package com.geedgenetworks.core.udf.cn; + +import com.geedgenetworks.common.Event; +import com.geedgenetworks.common.udf.UDFContext; +import org.apache.flink.api.common.functions.RuntimeContext; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +import static com.geedgenetworks.core.udf.cn.LookupTestUtils.*; +import static org.junit.jupiter.api.Assertions.assertEquals; + +class VpnLookupTest { + + private VpnLookup vpnLookup; + + @BeforeEach + void setUp() { + vpnLookup = new VpnLookup(); + } + + @Test + void evaluateIp() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "IP_TO_VPN"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("server_ip")); + udfContext.setOutput_fields(Collections.singletonList("server_vpn_service_name")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "addr_format,ip1,ip2,vpn_service_name\nCIDR,116.178.65.0,25,阿里"; + mockKnowledgeBaseHandler(content); + + vpnLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("server_ip", "116.178.65.100"); + event.setExtractedFields(fields); + Event evaluate = vpnLookup.evaluate(event); + assertEquals("阿里", evaluate.getExtractedFields().get("server_vpn_service_name")); + } + + @Test + void evaluateDomain() { + UDFContext udfContext = new UDFContext(); + Map<String, Object> parameters = new HashMap<>(); + parameters.put("kb_name", kbName); + parameters.put("option", "DOMAIN_TO_VPN"); + udfContext.setParameters(parameters); + udfContext.setLookup_fields(Collections.singletonList("domain")); + udfContext.setOutput_fields(Collections.singletonList("domain_vpn_service_name")); + + RuntimeContext runtimeContext = mockRuntimeContext(); + + String content = "domain,vpn_service_name\nblade43.singapore-rack443.nodes.gen4.ninja,cyberghostvpn"; + mockKnowledgeBaseHandler(content); + + vpnLookup.open(runtimeContext, udfContext); + + Event event = new Event(); + Map<String, Object> fields = new HashMap<>(); + fields.put("domain", "test.blade43.singapore-rack443.nodes.gen4.ninja"); + event.setExtractedFields(fields); + Event evaluate = vpnLookup.evaluate(event); + assertEquals("cyberghostvpn", evaluate.getExtractedFields().get("domain_vpn_service_name")); + } + + @AfterEach + void afterAll() { + clearState(); + } +}
\ No newline at end of file |
