def run(self):
print(banner)
dt = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
print(f'[*] Starting OneForAll @ {dt}\n')
logger.log('INFOR', f'开始运行OneForAll')
self.domains = utils.get_domains(self.target)
if self.domains:
for self.domain in self.domains:
collect = Collect(self.domain, export=False)
collect.run()
if self.brute:
# 由于爆破会有大量dns解析请求 并发常常会导致其他任务中的网络请求超时
brute = AIOBrute(self.domain)
brute.run()
table_name = self.domain.replace('.', '_')
db_conn = database.connect_db()
self.datas = database.get_data(db_conn, table_name).as_dict()
loop = asyncio.get_event_loop()
asyncio.set_event_loop(loop)
self.datas = loop.run_until_complete(resolve.bulk_query_a(self.datas))
self.datas = loop.run_until_complete(request.bulk_get_request(self.datas, self.port))
loop.run_until_complete(asyncio.sleep(0.25)) # 在关闭事件循环前加入一小段延迟让底层连接得到关闭的缓冲时间
loop.close()
database.clear_table(db_conn, table_name)
database.save_db(db_conn, table_name, self.datas)
# 数据库导出
if not self.path:
self.path = config.result_save_path.joinpath(f'{self.domain}.{self.format}')
dbexport.export(table_name, db_conn, self.valid, self.path, self.format, self.output)
db_conn.close()
else:
logger.log('FATAL', f'获取域名失败')
logger.log('INFOR', f'结束运行OneForAll')
def save_db(self):
lock.acquire()
db_conn = database.connect_db()
table_name = self.domain.replace('.', '_')
database.create_table(db_conn, table_name)
source, results = self.results
database.save_db(db_conn, table_name, results, source) # 将结果存入数据库中
lock.release()
def run(self, rx_queue=None):
self.domains = utils.get_domains(self.target)
while self.domains:
self.domain = self.domains.pop()
start = time.time()
db_conn = database.connect_db()
table_name = self.domain.replace('.', '_')
database.create_table(db_conn, table_name)
if not rx_queue:
rx_queue = queue.Queue()
logger.log('INFOR', f'开始执行{self.source}模块爆破域名{self.domain}')
logger.log(
'INFOR',
f'{self.source}模块使用{self.processes}个进程乘{self.coroutine}个协程')
# logger.log('INFOR', f'{self.source}模块使用个进程乘{self.coroutine}个协程')
if self.recursive_brute and not self.fuzz: # fuzz模式不使用递归爆破
logger.log('INFOR', f'开始递归爆破{self.domain}的第1层子域')
loop = asyncio.get_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(self.main(self.domain, rx_queue))
# 递归爆破下一层的子域
if self.recursive_brute and not self.fuzz: # fuzz模式不使用递归爆破
for layer_num in range(
1,
self.recursive_depth): # 之前已经做过1层子域爆破 当前实际递归层数是layer+1
logger.log('INFOR',
f'开始递归爆破{self.domain}的第{layer_num + 1}层子域')
for subdomain in self.subdomains.copy():
if subdomain.count('.') - self.domain.count(
'.') == layer_num: # 进行下一层子域爆破的限制条件
loop.run_until_complete(
self.main(subdomain, rx_queue))
while not rx_queue.empty(): # 队列不空就一直取数据存数据库
source, results = rx_queue.get()
database.save_db(db_conn, table_name, results,
source) # 将结果存入数据库中
database.copy_table(db_conn, table_name)
database.deduplicate_subdomain(db_conn, table_name)
database.remove_invalid(db_conn, table_name)
end = time.time()
self.elapsed = round(end - start, 1)
logger.log('INFOR', f'结束执行{self.source}模块爆破域名{self.domain}')
logger.log(
'INFOR',
f'{self.source}模块耗时{self.elapsed}秒发现{self.domain}的域名{len(self.subdomains)}个'
)
logger.log(
'DEBUG',
f'{self.source}模块发现{self.domain}的的域名 {self.subdomains}')
def export(table, db=None, valid=None, path=None, format='xlsx', output=False):
"""
OneForAll数据库导出模块
Example:
python dbexport.py --db result.db --table name --format csv --output False
python dbexport.py --db result.db --table name --format csv --path= ./result.csv
Note:
参数valid可选值1，0，None，分别表示导出有效，无效，全部子域
参数format可选格式：'csv', 'tsv', 'json', 'yaml', 'html', 'xls', 'xlsx', 'dbf', 'latex', 'ods'
参数path为None会根据format参数和域名名称在项目结果目录生成相应文件
:param str table: 要导出的表
:param str db: 要导出的数据库路径(默认为results/result.sqlite3)
:param int valid: 导出子域的有效性(默认None)
:param str format: 导出格式(默认xlsx)
:param str path: 导出路径(默认None)
:param bool output: 是否将导出数据输出到终端(默认False)
"""
db_conn = database.connect_db(db)
if valid is None:
rows = database.get_data(db_conn, table)
elif isinstance(valid, int):
rows = database.get_subdomain(db_conn, table, valid)
else:
rows = database.get_data(db_conn, table) # 意外情况导出全部子域
if output:
print(rows.dataset)
if not path:
path = 'export.' + format
logger.log('INFOR', f'正在将数据库中{table}表导出')
try:
with open(path, 'w') as file:
file.write(rows.export(format))
logger.log('INFOR', '成功完成导出')
logger.log('INFOR', path)
except TypeError:
with open(path, 'wb') as file:
file.write(rows.export(format))
logger.log('INFOR', '成功完成导出')
logger.log('INFOR', path)
except Exception as e:
logger.log('ERROR', e)
def run(self, rx_queue=None):
"""
类运行入口
"""
start = time.time()
self.get_mod()
self.import_func()
if not rx_queue:
rx_queue = queue.Queue(maxsize=len(self.collect_func)) # 结果集队列
threads = []
# 创建多个子域收集线程
for collect_func in self.collect_func:
thread = threading.Thread(target=collect_func,
args=(self.domain, rx_queue),
daemon=True)
threads.append(thread)
# 启动所有线程
for thread in threads:
thread.start()
# 等待所有线程完成
for thread in threads:
thread.join()
db_conn = database.connect_db()
table_name = self.domain.replace('.', '_')
database.create_table(db_conn, table_name)
database.copy_table(db_conn, table_name)
database.deduplicate_subdomain(db_conn, table_name)
database.remove_invalid(db_conn, table_name)
db_conn.close()
# 数据库导出
if self.export:
if not self.path:
self.path = config.result_save_path.joinpath(
f'{self.domain}.{self.format}')
dbexport.export(table_name, path=self.path, format=self.format)
end = time.time()
self.elapsed = round(end - start, 1)
from common.storage import get_storage_interface
from . import director_config
from .driver import create_driver_from_env
from .orchestrator import Orchestrator
from .terminator import Terminator
def _exit_gracefully(signum, frame):
orchestrator.shutdown()
log_util.configure_logging(__package__, logging.DEBUG)
signal.signal(signal.SIGINT, _exit_gracefully)
signal.signal(signal.SIGTERM, _exit_gracefully)
db_interface = Database(connect_db(
director_config).get_database()) # Docs say PyMongo is thread-safe
worker_driver = create_driver_from_env(director_config)
orchestrator = Orchestrator(db_interface, worker_driver)
terminator = Terminator(db_interface, worker_driver,
get_storage_interface(director_config))
# Run terminator (worker garbage collection) on a separate thread
thread_gc = threading.Thread(name='worker_gc', target=terminator.run)
thread_gc.start()
orchestrator.run()
# When orchestrator stops in the main thread, wait for terminator to end too
terminator.shutdown()
thread_gc.join()
import sys
from common import log_util
from common.database import Database, connect_db
from common.storage import get_storage_interface
from . import worker_config
from .worker import Worker
def _exit_gracefully(signum, frame):
worker.shutdown()
log_util.configure_logging(__package__, logging.DEBUG)
signal.signal(signal.SIGINT, _exit_gracefully)
signal.signal(signal.SIGTERM, _exit_gracefully)
_log = logging.getLogger('worker')
consumer_tag = os.environ.get('HARMONY_WORKER_ID')
if consumer_tag is None:
_log.critical('Attempted anonymous worker startup (no consumer tag given)')
sys.exit(1)
worker = Worker(consumer_tag=consumer_tag,
db_interface=Database(
connect_db(worker_config).get_database()),
storage_interface=get_storage_interface(worker_config))
_log.info('Started worker "%s"', consumer_tag)
worker.run()
def _connect_db():
return Database(connect_db(current_config).get_database())
from common.database import Database, connect_db
from apiserver.config import current_config
db = Database(connect_db(current_config).get_database())
print(db.get_library('5ece9573b16fafc9c3cbd2a4'))
arg_parser.add_argument('--clean',
action='store_true',
help="clean the database and storage before adding")
args = arg_parser.parse_args()
if not os.path.isdir(args.folder):
print(f'"{args.folder}" is not a folder', file=sys.stderr)
sys.exit(1)
lastfm_api_key = os.environ.get('LASTFM_API_KEY')
if lastfm_api_key is None:
print(
'LASTFM_API_KEY is not set (you can scrape for some on Google), tags and artist bio will not be loaded',
file=sys.stderr)
db = Database(connect_db(config).get_database())
st = get_storage_interface(config)
if args.clean:
db.artists.drop()
st.delete_all_files(config.STORAGE_BUCKET_REFERENCE)
st.delete_all_files(config.STORAGE_BUCKET_TRANSCODED)
st.delete_all_files(config.STORAGE_BUCKET_IMAGES)
print('Cleaned up database and storage', file=sys.stderr)
songs = [
os.path.join(dp, f) for dp, dn, filenames in os.walk(args.folder)
for f in filenames if os.path.splitext(f)[1].lower() == '.flac'
]
meta_tree = make_tree(songs)
insert_data(meta_tree)
