Get IP from buffer on geoip too

This commit is contained in:
JeDaYoshi 2021-10-17 21:42:15 -04:00
parent e68af22d60
commit 2e5836ae22
No known key found for this signature in database
GPG key ID: 8060B288C274219D

View file

@ -21,6 +21,17 @@ def _parse(value):
@utils.export("channelset", utils.FunctionSetting(_parse, "dns-nameserver", @utils.export("channelset", utils.FunctionSetting(_parse, "dns-nameserver",
"Set DNS nameserver", example="8.8.8.8")) "Set DNS nameserver", example="8.8.8.8"))
class Module(ModuleManager.BaseModule): class Module(ModuleManager.BaseModule):
def _get_ip(self, event):
ip = event["args_split"][0] if event["args"] else ""
if not ip:
line = event["target"].buffer.find(REGEX_IP)
if line:
ip = line.match
if not ip:
raise utils.EventError("No IP provided")
return ip
@utils.hook("received.command.dig", alias_of="dns") @utils.hook("received.command.dig", alias_of="dns")
@utils.hook("received.command.dns", min_args=1) @utils.hook("received.command.dns", min_args=1)
def dns(self, event): def dns(self, event):
@ -79,14 +90,16 @@ class Module(ModuleManager.BaseModule):
(t, ttl, ", ".join(r)) for t, ttl, r in results] (t, ttl, ", ".join(r)) for t, ttl, r in results]
event["stdout"].write("(%s) %s" % (hostname, " | ".join(results_str))) event["stdout"].write("(%s) %s" % (hostname, " | ".join(results_str)))
@utils.hook("received.command.geoip", min_args=1) @utils.hook("received.command.geoip")
def geoip(self, event): def geoip(self, event):
""" """
:help: Get geoip data on a given IPv4/IPv6 address :help: Get GeoIP data on a given IPv4/IPv6 address
:usage: <IP> :usage: <IP>
:prefix: GeoIP :prefix: GeoIP
""" """
page = utils.http.request(URL_GEOIP % event["args_split"][0]).json() ip = self._get_ip(event)
page = utils.http.request(URL_GEOIP % ip).json()
if page: if page:
if page["status"] == "success": if page["status"] == "success":
try: try:
@ -117,13 +130,7 @@ class Module(ModuleManager.BaseModule):
:usage: <IP> :usage: <IP>
:prefix: rDNS :prefix: rDNS
""" """
ip = event["args_split"][0] if event["args"] else "" ip = self._get_ip(event)
if not ip:
line = event["target"].buffer.find(REGEX_IP)
if line:
ip = line.match
if not ip:
raise utils.EventError("No IP provided")
try: try:
hostname, alias, ips = socket.gethostbyaddr(ip) hostname, alias, ips = socket.gethostbyaddr(ip)