This repository has been archived by the owner on May 31, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 7
/
app.py
77 lines (62 loc) · 1.74 KB
/
app.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
#!/usr/bin/python
# -*- coding: utf-8 -*-
import hashlib
import os
import sys
import asyncio
import click
from requests import get
from malwarefeeds.engine import Engine
SAMPLES_PATH = os.path.expanduser('~/.samples')
def store_file(content):
fhash = hashlib.sha256(content).hexdigest()
directory = os.path.join(
SAMPLES_PATH,
fhash[0],
fhash[1],
fhash[2],
fhash[3])
if not os.path.isdir(directory):
os.makedirs(directory)
filepath = os.path.join(directory, fhash)
if not os.path.exists(filepath):
with open(filepath, 'wb') as fd:
fd.write(content)
async def download_from_feeds(q):
e = Engine()
e.update()
for _, url in e.read():
await q.put(url)
await q.put(None)
async def download_samples(q):
while True:
url = await q.get()
if url is None:
break
print("Downloading URL: %s" % url)
r = get(url, stream=True)
store_file(r.content)
q.task_done()
@click.command()
@click.option('--repo', default=SAMPLES_PATH, help='The directory to store files')
def cli(repo):
try:
if not os.path.isdir(repo):
print("Creating directory %s" % repo)
os.makedirs(repo)
SAMPLES_PATH = repo
print("Storing files in %s" % SAMPLES_PATH)
ioloop = asyncio.get_event_loop()
q = asyncio.Queue(loop=ioloop)
producer = download_from_feeds(q)
consumer = download_samples(q)
ioloop.run_until_complete(asyncio.gather(producer, consumer))
ioloop.close()
except KeyboardInterrupt:
pass
except PermissionError as ex:
print(ex)
except SystemExit:
pass
if __name__ == "__main__":
cli(None)