[백신 개발](4)[바이러스 db]

바이러스 db 구축

아무래도 새로운 악성코드가 나타날 때 마다 vaccine.py를 수정하는 건 무리다.
따라서 vaccine파일과 db를 분리하는 것이 좋을 것이다.

 

csv 파일과 pandas dataframe을 가지고 db를 구축했다.

csv를 읽는건 pandas.read_csv(<file>)로 처리하면 된다.

import sys
import hashlib
import os
import pandas as pd

# EICAR_new.txt 파일 읽기
fp = open('test.txt', 'rb')
fbuf = fp.read()
fp.close()

malware_db = pd.read_csv("malware.csv")
# MD5 해시를 이용한 진단
input_hash = hashlib.md5(fbuf).hexdigest()
if input_hash in malware_db['hash']:
    print("virus!")
else:
    print("benign!")

그리고 이 db 파일을 누군가 열어 맘대로 수정하면 바이러스 탐지가 제대로 안될 수 있다.

따라서 약간의 db파일 암호화도 필요하겠다.

import sys
import hashlib
import os
import zlib

# cmd 명령어로 KAVM형식으로 암호화해준다
def main():
    # 오류: 인자가 없음
    if len(sys.argv) != 2 :
        print('Usage : kmake.py [file]')
        return

    # 파일 읽기
    fname = sys.argv[1]
    tname = fname
    fp = open(tname, 'rb') 
    buf = fp.read()
    fp.close()

    # 파일 압축
    buf2 = zlib.compress(buf)  

    # 압축 내용 0xFF로 XOR
    buf3 = b'' 
    for c in buf2 : 
        buf3 += (c ^ 0xFF).to_bytes(1, byteorder="little") 

    # 헤더 붙이기
    buf4 = b'KAVM' + buf3 

    # md5로 돌린 내용을 암호화 파일 뒤에 붙임
    # (파일 손상 검증용)
    f = buf4
    for i in range(3) : 
        md5 = hashlib.md5()
        md5.update(f)
        f = md5.hexdigest()
        f = bytes(f, 'utf-8') 
    buf4 += f 

    # .kmd 파일로 만든다
    kmd_name = fname.split('.')[0] + '.kmd' 
    fp = open(kmd_name, 'wb') # ❻ kmd 확장자로 암호 파일을 만든다. 
    fp.write(buf4)
    fp.close()

    print ('%s -> %s' % (fname, kmd_name)) # 결과를 출력한다. 

if __name__ == '__main__' :
    main()

암호화는 책에서 나온 것처럼 단순하게

1. 0xFF XOR 암호화
2. md5검증
   이 2가지를 수행하도록 했다.

# KMD 파일을 복호화한다. 
def DecodeKMD(fname):
    try:
        fp = open(fname, 'rb') 
        buf = fp.read()
        fp.close()

        buf2 = buf[:-32] # 암호화 내용을 분리한다. 
        fmd5 = buf[-32:] # MD5를 분리한다. 
        f = buf2

        # md5 비교
        for i in range(3) : # 암호화 내용의 MD5를 구한다. 
            md5 = hashlib.md5()
        md5.update(f)
        f = md5.hexdigest()
        f = bytes(f, 'utf-8')
        if f != fmd5 : 
            raise SystemError

        # 0xFF 복호화
        buf3 = b'' 
        for c in buf2[4:] : 
            buf3 += (c ^ 0xFF).to_bytes(1, byteorder="little") 

        # 압축 해제
        buf4 = zlib.decompress(buf3) # ❸ 압축을 해제한다. 
        return buf4 # 성공했다면 복호화된 내용을 리턴한다. 
    except :
        pass
    return None

물론 이걸 해석하는 로직도 백신에 넣어줘야 한다.

 

나는 책과 달리 csv, pandas로 db를 관리할거라 적절히 이식해주면...

import sys
import hashlib
import zlib
import os
import pandas as pd
import io

def main():
    # EICAR_new.txt 파일 읽기
    fp = open('test.txt', 'rb')
    fbuf = fp.read()
    fp.close()

    db = DecodeKMD("malware.kmd")
    malware_db = pd.read_csv(io.BytesIO(db))
    # MD5 해시를 이용한 진단
    input_hash = hashlib.md5(fbuf).hexdigest()
    if input_hash in malware_db['hash']:
        print("virus!")
        # 치료하기
        #os.remove(malware_hash.get(input_hash)) # 삭제로 치료는 주의가 필요
    else:
        print("benign!")


# KMD 파일을 복호화한다. 
def DecodeKMD(fname):
    try:
        fp = open(fname, 'rb') 
        buf = fp.read()
        fp.close()

        buf2 = buf[:-32] # 암호화 내용을 분리한다. 
        fmd5 = buf[-32:] # MD5를 분리한다. 
        f = buf2

        # md5 비교
        for _ in range(3) : # 암호화 내용의 MD5를 구한다. 
            md5 = hashlib.md5()
            md5.update(f)
            f = md5.hexdigest()
            f = bytes(f, 'utf-8')

        if f != fmd5 : 
            raise SystemError

        # 0xFF 복호화
        buf3 = b'' 
        for c in buf2[4:] : 
            buf3 += (c ^ 0xFF).to_bytes(1, byteorder="little") 

        # 압축 해제
        buf4 = zlib.decompress(buf3) # ❸ 압축을 해제한다. 
        return buf4 # 성공했다면 복호화된 내용을 리턴한다. 
    except :
        pass
    return None

if __name__ == "__main__":
    main()

잘 작동한다!

 

모든 자료는 여기에
https://github.com/VeryBigsilver/malware-vaccine/tree/main/3_malware_db