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

78 lines
2.2 KiB
Python
Raw Normal View History

#!/usr/bin/env python3
from datetime import datetime
2019-02-19 04:40:35 +00:00
from hashlib import sha1
from pathlib import Path
from sqlalchemy import LargeBinary, Column, ForeignKey, Integer, String, DateTime, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, relationship, backref
from sqlalchemy.orm.collections import attribute_mapped_collection
import json
from .disassemble import disasm
now = '{0:%Y%m%dT%H%M%S}'.format(datetime.utcnow())
Base = declarative_base()
2019-02-19 23:04:45 +00:00
2019-02-19 04:40:35 +00:00
def db_config(path: Path) -> Session:
engine = create_engine(f'sqlite:///{path.resolve()}', native_datetime=True)
Base.metadata.create_all(engine)
session = Session(engine)
return session
2019-02-19 23:04:45 +00:00
2019-02-19 04:40:35 +00:00
class ScrapNode(Base):
__tablename__ = 'scrap_node'
ctime = Column(DateTime, default=datetime.utcnow)
id = Column(Integer, primary_key=True)
length = Column(Integer, default=0)
mtime = Column(DateTime, onupdate=datetime.utcnow)
parent_id = Column(Integer, ForeignKey(id))
checksum = Column(String)
disasm = Column(String)
2019-02-19 04:40:35 +00:00
image = Column(LargeBinary)
children = relationship(
2019-02-19 04:40:35 +00:00
"ScrapNode",
cascade="all, delete-orphan",
backref=backref("parent", remote_side=id),
collection_class=attribute_mapped_collection('name'))
2019-02-19 04:40:35 +00:00
def __init__(self, *, child: bytes, parent_id: int = None):
self.parent_id = parent_id
self.image = child
self.length = len(child)
2019-02-19 04:40:35 +00:00
self.sha1sum
self.disasm = disasm(child)
2019-02-19 04:40:35 +00:00
def __repr__(self):
values = {
'checksum': self.checksum,
'length': self.length,
'disasm': self.disasm,
2019-02-19 04:40:35 +00:00
'parent_id': self.parent_id,
'id': self.id,
2019-02-19 04:40:35 +00:00
}
return json.dumps(values, indent=1)
@property
def sha1sum(self):
if self.checksum:
return self.checksum
checksum = sha1()
checksum.update(self.image)
self.checksum = checksum.hexdigest()
return self.checksum
def disasm(shellcode: bytes) -> str:
opcodes = list()
for opcode in capstone.disasm(shellcode, 0):
opcodes += f'{opcode.mnemonic} {opcode.op_str}\n'
return opcodes