jkhsjdhjs/nxy
1
0
Fork 0
Code Issues 5 Pull Requests Projects Releases Wiki Activity

urlinfo: limit download size + check content type

Browse Source
This commit is contained in:
jkhsjdhjs 2021-05-14 22:13:30 +00:00
parent 233cfc9089
commit c4f66e5346
1 changed files with 19 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
bot/urlinfo.py
View File

@ -15,16 +15,32 @@ class URLInfo(Plugin):
"^https?:\/\/f0ck\.me" "^https?:\/\/f0ck\.me"
] ]
SIZE_LIMIT = 2 * 1024 ** 2
@irc3.event(r'(?i)^:\S+ PRIVMSG (?P<target>\S+) :.*(?P<url>https?:\/\/\S+\.\S+).*') @irc3.event(r'(?i)^:\S+ PRIVMSG (?P<target>\S+) :.*(?P<url>https?:\/\/\S+\.\S+).*')
def url_parser(self, target: str, url: str): def url_parser(self, target: str, url: str):
for regex in self.BLACKLIST: for regex in self.BLACKLIST:
if re.match(regex, url): if re.match(regex, url):
return return
bytes_io = io.BytesIO()
try: try:
response = requests.get(url, timeout=10) with requests.get(url, timeout=10, stream=True) as r:
except requests.exceptions.ConnectionError: r.raise_for_status()
size = 0
for chunk in r.iter_content(chunk_size=1024 ** 2):
size += len(chunk)
if size >= self.SIZE_LIMIT:
return
bytes_io.write(chunk)
except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError) as e:
return return
tree = etree.parse(io.StringIO(response.text), etree.HTMLParser()).getroot()
mime_type = r.headers.get("content-type")
if mime_type is not None:
if mime_type.split(";")[0] != "text/html":
return
bytes_io.seek(0)
tree = etree.parse(bytes_io, etree.HTMLParser()).getroot()
title = tree.xpath("/html/head/title") title = tree.xpath("/html/head/title")
if len(title) > 0: if len(title) > 0:
self.bot.privmsg(target, '\x02[URLInfo]\x02 ' + title[0].text.strip()) self.bot.privmsg(target, '\x02[URLInfo]\x02 ' + title[0].text.strip())