VTZero.py (3130B)
1 from flask import Flask, request, render_template_string
2 import asyncio
3 import vt
4 import County_Codes
5 from dotenv import load_dotenv
6 import os
7
8 # .env 파일이 존재하지 않는 경우 생성
9 if not os.path.exists('.env'):
10 api_key = input("Please enter your VirusTotal API key: ")
11 with open('.env', 'w') as f:
12 f.write(f"API_KEY={api_key}\n")
13 print(".env file created successfully.")
14 else:
15 print(".env file already exists.")
16
17 # .env 파일에서 환경 변수 로드
18 load_dotenv()
19
20 API_KEY = os.getenv('API_KEY') # .env 파일에서 API 키 가져오기
21
22 if not API_KEY:
23 raise ValueError("API_KEY is not set. Please set the API_KEY in the .env file.")
24
25 app = Flask(__name__)
26
27 TEMPLATE = '''
28 <!doctype html>
29 <html>
30 <head><title>IP Information</title></head>
31 <body>
32 <h1>IP Information</h1>
33 <form method="post">
34 <textarea name="ips" rows="10" cols="50" placeholder="Enter IPs, one per line"></textarea><br>
35 <input type="submit" value="Submit">
36 </form>
37 {% if results %}
38 <h2>Results</h2>
39 <table border="1">
40 <tr>
41 <th>IP</th>
42 <th>국가</th>
43 <th>소유자</th>
44 <th>유해</th>
45 </tr>
46 {% for ip, info in results.items() %}
47 <tr>
48 <td>{{ ip }}</td>
49 <td>{{ info.country }}</td>
50 <td>{{ info.as_owner }}</td>
51 <td style="color: {{ 'red' if info.malicious|int >= 1 else 'black' }}">{{ info.malicious }}</td>
52 </tr>
53 {% endfor %}
54 </table>
55 {% endif %}
56 </body>
57 </html>
58 '''
59
60 async def get_ip_info(client, ip):
61 result = {'country': '', 'as_owner': '', 'malicious': ''}
62 try:
63 ip_info = await client.get_object_async(f"/ip_addresses/{ip}")
64 ip_info_dict = ip_info.to_dict() # Object를 dictionary로 변환
65
66 if 'attributes' in ip_info_dict:
67 attributes = ip_info_dict['attributes']
68 if 'country' in attributes:
69 country_code = attributes['country']
70 result['country'] = County_Codes.country_code_to_korean.get(country_code, country_code)
71 if 'last_analysis_stats' in attributes:
72 result['malicious'] = str(attributes['last_analysis_stats']['malicious'])
73 if 'as_owner' in attributes:
74 result['as_owner'] = attributes['as_owner']
75 except Exception as e:
76 print(f"Error fetching data for IP {ip}: {e}")
77 return result
78
79 @app.route('/', methods=['GET', 'POST'])
80 def index():
81 results = {}
82 if request.method == 'POST':
83 ips = request.form['ips'].split()
84 loop = asyncio.new_event_loop()
85 asyncio.set_event_loop(loop)
86 client = vt.Client(API_KEY)
87 tasks = [get_ip_info(client, ip) for ip in ips]
88 results_list = loop.run_until_complete(asyncio.gather(*tasks))
89 for ip, info in zip(ips, results_list):
90 results[ip] = info
91 client.close()
92 return render_template_string(TEMPLATE, results=results)
93
94 if __name__ == '__main__':
95 app.run(host='0.0.0.0', port=5000, debug=True) # 모든 IP 주소에서 접근 가능