scrap-is-not-scrap/sins/run.py

119 lines
3.4 KiB
Python
Raw Normal View History

#! /usr/bin/env python3
from argparse import ArgumentParser
from datetime import datetime
from multiprocessing import Process, Queue
2019-02-19 03:33:56 +00:00
from pathlib import Path
from queue import Empty
2019-02-19 04:40:35 +00:00
from sqlalchemy import exists, desc
from tempfile import TemporaryDirectory
import logging
from .mutation import generation, mutate, seed_shell, growth
2019-02-19 04:40:35 +00:00
from .orm import db_config, ScrapNode
2019-02-18 23:16:25 +00:00
def sins():
now = '{0:%Y%m%dT%H%M%S}'.format(datetime.utcnow())
parser = ArgumentParser(
description='position independent code (PIC) mutation experiment.')
parser.add_argument('-v', '--verbose', action='count')
2019-02-18 23:16:25 +00:00
parser.add_argument('-s', '--seed', help='path to PIC image.')
parser.add_argument('-o', '--output', help='path to results directory.')
parser.add_argument('-l', '--lineage', default=10,
2019-02-19 03:33:56 +00:00
help='max count of unsuccessful generation.')
args = parser.parse_args()
log_level = logging.INFO
log_format = logging.Formatter('%(message)s')
if args.verbose:
log_level = logging.DEBUG
log_format = logging.Formatter(
'%(levelname)s %(filename)s:%(lineno)d\n%(message)s\n')
logger = logging.getLogger('sins')
logger.setLevel(log_level)
stream_handler = logging.StreamHandler()
stream_handler.setLevel(log_level)
stream_handler.setFormatter(log_format)
logger.addHandler(stream_handler)
if args.output:
log_path = f'{args.output}/sins-{now}.log'
file_handler = logging.FileHandler(log_path)
file_handler.setLevel(log_level)
file_handler.setFormatter(log_format)
logger.addHandler(file_handler)
2019-02-18 23:16:25 +00:00
logger.info(now)
2019-02-19 04:40:35 +00:00
if args.output:
db_path = Path(f'{args.output}/sins.sqlite')
else:
temp_dir = TemporaryDirectory()
db_path = Path(f'{temp_dir.name}/sins.sqlite')
session = db_config(db_path)
logger.info(f'db_path: {db_path}')
recent = session.query(ScrapNode).order_by(desc('ctime')).first()
2019-02-19 04:40:35 +00:00
if args.seed:
seed_path = Path(args.seed)
with seed_path.open('rb') as seed_file:
seed_data = seed_file.read()
seed = ScrapNode(child=seed_data)
2019-02-19 23:04:45 +00:00
exists = session.query(ScrapNode).filter(
ScrapNode.checksum == seed.checksum)
2019-02-19 04:40:35 +00:00
if exists:
seed = exists[0]
else:
session.add(seed)
session.commit()
logger.debug(f'args.seed:\n{seed}')
elif recent:
seed = recent
logger.debug(f'recent:\n{seed}')
2019-02-19 04:40:35 +00:00
else:
seed = ScrapNode(child=seed_shell)
session.add(seed)
session.commit()
logger.debug(f'seed_shell:\n{seed}')
2019-02-19 04:40:35 +00:00
parent = seed
queue = Queue()
while True:
lineage = 0
while lineage < args.lineage:
scrap = mutate(parent.image)
logger.debug(f'lineage: {lineage}')
result = None
2019-02-19 00:28:34 +00:00
2019-02-19 04:40:35 +00:00
proc = Process(target=generation, args=(queue, scrap))
proc.start()
try:
result = queue.get(timeout=1)
except Empty:
lineage += 1
continue
if not result:
lineage += 1
continue
2019-02-18 23:16:25 +00:00
scrap = growth(shellcode=scrap, length=result)
2019-02-19 04:40:35 +00:00
parent = ScrapNode(child=scrap, parent_id=parent.id)
session.add(parent)
2019-02-19 04:40:35 +00:00
session.commit()
2019-02-19 04:40:35 +00:00
logger.info(f'scrap:\n{parent}')
lineage = 0