File size: 4,648 Bytes
7ac32ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Run termite-classifier with PyTorch. It reads plain data files and never runs a command.
Usage: python termite_classifier.py 'ls -la' [FOLDER] [SHELL] [OS] [SESSION]"""
import json, pathlib, struct, sys
import torch, torch.nn as nn, torch.nn.functional as F
HERE = pathlib.Path(__file__).resolve().parent
QUESTIONS = ("only_reads", "no_secrets", "harmful")
TYPES = {"F32": torch.float32}
CLS, FIRST = 1, 3
def pieces(data):
"""A piece is a run of ASCII letters, digits or bytes over 127, or one other byte."""
start = 0
for i, b in enumerate(data):
if not (48 <= b <= 57 or 65 <= b <= 90 or 97 <= b <= 122 or b > 127):
if i > start:
yield data[start:i]
yield data[i:i + 1]
start = i + 1
if start < len(data):
yield data[start:]
class Tokeniser:
def __init__(self, merges):
self.rank = {tuple(m): i for i, m in enumerate(merges)}
def piece(self, data):
w = [FIRST + b for b in data]
while len(w) > 1:
found = min((self.rank.get(p, 1 << 30), i) for i, p in enumerate(zip(w, w[1:])))
if found[0] == 1 << 30:
break
w[found[1]:found[1] + 2] = [FIRST + 256 + found[0]]
return w
def encode(self, text):
ids = [CLS]
for p in pieces(text.encode("utf8", "surrogateescape")):
ids += self.piece(p)
return ids
class Block(nn.Module):
def __init__(self, width, heads):
super().__init__()
self.heads = heads
self.n1, self.n2 = nn.LayerNorm(width), nn.LayerNorm(width)
self.qkv, self.out = nn.Linear(width, 3 * width), nn.Linear(width, width)
self.up, self.down = nn.Linear(width, 4 * width), nn.Linear(4 * width, width)
def forward(self, x, keep):
b, t, w = x.shape
q, k, v = self.qkv(self.n1(x)).view(b, t, 3, self.heads, w // self.heads).permute(2, 0, 3, 1, 4)
a = F.scaled_dot_product_attention(q, k, v, attn_mask=keep[:, None, None, :])
x = x + self.out(a.transpose(1, 2).reshape(b, t, w))
return x + self.down(F.gelu(self.up(self.n2(x))))
class Encoder(nn.Module):
def __init__(self, vocab, longest, layers, width, heads):
super().__init__()
self.tokens, self.places = nn.Embedding(vocab, width), nn.Embedding(longest, width)
self.blocks = nn.ModuleList(Block(width, heads) for _ in range(layers))
self.last = nn.LayerNorm(width)
self.heads = nn.Linear(width, 3)
def forward(self, ids):
keep = ids != 0
x = self.tokens(ids) + self.places(torch.arange(ids.shape[1]))
for block in self.blocks:
x = block(x, keep)
return self.heads(self.last(x)[:, 0])
def read_weights(path):
with open(path, "rb") as f:
size = struct.unpack("<Q", f.read(8))[0]
header = json.loads(f.read(size))
body = bytearray(f.read())
return {name: torch.frombuffer(body, dtype=TYPES[t["dtype"]], offset=t["data_offsets"][0],
count=(t["data_offsets"][1] - t["data_offsets"][0]) // 4).view(t["shape"])
for name, t in header.items() if name != "__metadata__"}
class Classifier:
def __init__(self, folder=HERE):
folder = pathlib.Path(folder)
self.config = json.load(open(folder / "config.json"))
c = self.config
self.tok = Tokeniser(json.load(open(folder / "merges.json")))
self.net = Encoder(c["vocab"], c["longest"], c["layers"], c["width"], c["heads"])
self.net.load_state_dict(read_weights(folder / c["weights"]))
self.net.eval()
@torch.no_grad()
def answer(self, text):
"""Chances for each question. A command too long to read gets the strict answer."""
ids = self.tok.encode(text)
if len(ids) > self.config["longest"]:
return {"only_reads": 0.0, "no_secrets": 0.0, "harmful": 1.0}
z = self.net(torch.tensor([ids]))[0]
return {q: torch.sigmoid(z[i] / self.config["temperatures"][i]).item() for i, q in enumerate(QUESTIONS)}
def safe(self, a):
limit = self.config["limits"]
return a["only_reads"] >= limit["only_reads"] and a["no_secrets"] >= limit["no_secrets"] and a["harmful"] <= limit["harmful"]
if __name__ == "__main__":
command, *rest = sys.argv[1:]
folder, shell, system, session = (rest + ["/home/dev/src/project", "bash", "linux", "local"][len(rest):])[:4]
model = Classifier()
answer = model.answer(f"{shell} {system} {session}\n{folder}\n{command}")
print(json.dumps({**answer, "safe": model.safe(answer)}, indent=1))
|