From ef01d0cdf53f0b1f9e8a25f21b1790d422688358 Mon Sep 17 00:00:00 2001 From: Antonino Lorenzo <94693967+antoninoLorenzo@users.noreply.github.com> Date: Thu, 20 Jun 2024 19:58:40 +0200 Subject: [PATCH] Planning Prompt Experiments --- src/__init__.py | 44 +++++ .../__pycache__/__init__.cpython-311.pyc | Bin 221 -> 270 bytes src/agent/__pycache__/agent.cpython-311.pyc | Bin 5002 -> 5727 bytes src/agent/__pycache__/llm.cpython-311.pyc | Bin 2146 -> 2148 bytes src/agent/__pycache__/plan.cpython-311.pyc | Bin 2679 -> 2704 bytes src/agent/__pycache__/prompts.cpython-311.pyc | Bin 2891 -> 8152 bytes src/agent/agent.py | 43 +++-- .../__pycache__/collections.cpython-311.pyc | Bin 4114 -> 4062 bytes .../knowledge/__pycache__/nlp.cpython-311.pyc | Bin 4422 -> 4433 bytes .../__pycache__/routing.cpython-311.pyc | Bin 5670 -> 5778 bytes .../__pycache__/store.cpython-311.pyc | Bin 10919 -> 10886 bytes src/agent/knowledge/store.py | 1 - src/agent/llm.py | 2 +- src/agent/plan.py | 1 + src/agent/prompts.py | 160 +++++++++++++++--- src/main.py | 36 ++-- 16 files changed, 225 insertions(+), 62 deletions(-) diff --git a/src/__init__.py b/src/__init__.py index e69de29..9123d84 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -0,0 +1,44 @@ +import json +from pathlib import Path + +from src.agent.knowledge import Store +from src.agent.knowledge import Collection, Document, Topic + + +def upload_knowledge(path: str, vdb: Store): + """Used to initialize and keep updated the Knowledge Base. + Already existing Collections will not be overwritten. + :param path: where the JSON datasets are located. + :param vdb: the reference to the Knowledge Base""" + base_path = Path(path) + + for i, p in enumerate(base_path.iterdir()): + if not (p.is_file() and p.suffix == '.json'): + continue + + if p.name in ['hack_tricks.json', 'null_byte.json']: + continue + + with open(str(p), 'r', encoding='utf-8') as file: + data = json.load(file) + + documents = [] + topics = set() + for item in data: + topic = Topic(item['category']) + topics.add(topic) + + document = Document( + name=item['title'], + content=item['content'], + topic=topic + ) + documents.append(document) + + collection = Collection( + collection_id=i, + title=p.name, + documents=documents, + topics=list(topics) + ) + vdb.create_collection(collection) diff --git a/src/agent/__pycache__/__init__.cpython-311.pyc b/src/agent/__pycache__/__init__.cpython-311.pyc index 45c046d3d10f7f3576e164a9d66205d25e3b9fab..055723e2d4daab027b275ba04593707eda05fe74 100644 GIT binary patch delta 142 zcmcc1*vG`XoR^o20SI=ul%%;&Mkh#0O?ZM#dWqN*5TECLWMvVFwBU0O+$LE&u=k delta 93 zcmeBUy35GBoR^o20SI`t3eprO@=9t70XfqdQW#Pga~N_NqZk<(Qka4nG?`z5RB1Ba nVs%VU%`5TKWSUqNFT?>9V+7)2Ss?L&nURrk;zvmqHXsK8HX0D@ diff --git a/src/agent/__pycache__/agent.cpython-311.pyc b/src/agent/__pycache__/agent.cpython-311.pyc index 05d245ffd7b0c2a83612f6282566b0ce676c82ab..a6f906744bd5c726916cd0220bafa4c548a9994b 100644 GIT binary patch literal 5727 zcmcH-U2oggl@v*dl4Z-XY$r~ln2nterM5T3O`4=lTDNwcq_&gB?z*a87)HEyl<32~ zq@6@?1`i7cVaP*QWUz*1fL@llXa?kQkNF9bK_CKw00D*p>(gAMPkq`sml7q$$<_`S zMID}d&$*ZPoO`}5|I*$b;^6xAoVoC?FvtBJ8;whe`m=##AS)+j)+c3x*-$dXb3Sg8)BSfiJ@CNCabLsVxsq)ptjpI~ zo6f_G!AyI$BiVr?hzS3V2f+AsgPnUAdH`$r8vf3ei~vuY-lm61nEz zWR!H4yBa;MPd!4S-nEQ#zlbDbFkd%`>EVZ+UMkPs?L?Y)Js@Rtmi-f5h0*u(Ia5pL z48^=klvGAD4CV6JBFUMThn%iFB}#sQ;l?e)B-xQ~a7_u`7N(}A;UUeMz>^Sc|1`LbRQ9Q0sbP9 z7W}$&M@;(lm>wVjT_)m!s0Z&zNm&ow;gUh@UGV8`z!R$DVZht!cst;rhZZ=FB z275OWxLm_o6`m9~({ zWGO1n7P2I33oH5+yHnLxv~)(hk|8SFY>E=v5obb`%>)i(x8uysriTu1YNbOnA=;ur zG7D^t6nBWW12Aw86}Y}XGJbM?oEX>5{L=gq%`eVhp!sVgWg7EiXNS*Un4LFhY972S zXF70X>6R_4YC4xTRkhH;Y9{x3Hctp1~qu3G)?VDFx%@xk)P_b+XHV8urt0eB00 zchz|-Bkx%K?>^pX^`C2Ht_+>B;-~+%V#PmT+!c8j+|OZ6+?*{KY26Mm1*UTzzd=hJ zfdBJB2)G||JO*|Cwzw+Lo>i_4gzeX=&+ARV&0TZ-sq6;w){*Ej*P7lWKJ@NHK`1T&CLM zsX4k#E{D+j@PaWJ0`UPkFG37A zb^sjRBsWbvkgcagrzdC5Bs%P%N}Xv5-3fbO`*F>hNqeC|yRoO;z|^YN2UAR$6dl4A zKc1bqXtymAQ#D{_h+uY4oy&kArRwRFsb(Nk(HOEzObd2!Yhoh_|Y!w0PJfo~-t82XL_ zu*HDPb?>}){7LNKX6)dGP>l^)v7sli_cmkiRb%5;Y`iGL8@qdpa;>ND-p|)BeR-*T zrrI-X^$c%uzTg`S6ent3J4!*TYhU?<)iqoR*Y8@myCjv;*3Ls4eb&y=YIwAGrg-Mr zOO`hSlZJ<_@bHVva~DeYf}T>XF$p6TMxy7q-5Bk_rhSEj6E{g}*(6S(26dXLVCGSn zh3HB;r$EAj>TBjJ?kit&a}AuizDIVLICHQhRp@dytF`Y4FdO*VHTK)SHZ(rAwiQ6u zF%Cm2a(IOq{!ZHKQra;zk6oM|)dXo-Kl z3ZL*fvJb)h+vKF!+mNp%c$|Q4105FRi^T5Ud=eShj0}|DuSO18kwX>v(0>r*Bn>1lM|0!ytTa_WaCDT`W6n0CQufw##=K<*OTZE9WT`5IxP^ z!1$&dFQuyTK1<$LY25!$G2`PTLmK|OO;A_B4oFGB_G%H|MG?Cl6ZP_rZ4Pd zRe8{o2P=)62?>gH@`jqzvLuBj*|^#c-tM_baDM0X_~Pa^ zWu}a)A#*}##xKf*!sWP`=6z^16yO~=0Te?4;O?@cfbx7*K4Qs7Dvf*lK4V~FLjML3 zR$^j72)J{-D!6%lla)&Xl#;?VH=ws3vi`{Ff``^pN&!v!lgP>?ACyZzR>v=yYm~FlX%GOQ+^%lOC&-t0TP@n*Ge)ao0> zqhzGZPVQG>trM5VH zZWzw<@r3hy3EZ<&g8BqNDaIaU7}Dr!!QSAdW1y4<@cX1`z67lJd zG*A|>ci=?;r2S>qyB|rd07#<_Y4p)75=UPIvR~R!j%{)9cqMRDi3-Fm%)(bDRfVrX zLHKNRF%P8-g@&IptROFvY*ssY^rFH*Z|kOhtolQQ16{KB;^*3twfl zW%z29Gt(*OBTd>cmLZcsEp&riG13rh-}wfs5_pe5a{5f7hfd>a+TrLtqT$hb1kd~E z_jat==YGXuc4T)pDNc#SJ`l9j711VxQ#u>$3js47dJZSw-Q;j5A3@0>d()4QBe()b z`#e?QdxOMQRzFU#Qra=#H;_+SqG-`xOZ{Gc+i{E2^b#ZjXP1{PSm(~ z#e3Jd{T1(BXUvf{Pr?2;4& zg+mUgN)G8r1*xq>d!UIz+Z=nSl=Rk9WmlGBtrYdz-W;W;wrcxkZAhTa?9MlD-g`6i z{$_sJ`%#zjRWRs7(B6H+o_O7NS7}=4zz8iP9qBlWOsrwdXfDf{yvCb?CYYiox;#EB znI6q!%9?C?HLuGPvOd$V`7z>9AJWBJNSD?)gdRYzHZ5QTb3*YbJ+Q`u40=r#)`MEa2(L!yX{kZb6(e%@ z?>O&Y578aC>v-J=(xI#2-Db_P4mxiL6Ou065;Rd44QYbYC7?Yu+5@!AXk9k= z2|@S5-0RYu?gQFa)AIxEuh9XZ1MZrjyGDU|Q1^wv_&P^CJ!#Ne+%fJOM1SC}Njf)& z#zguB-;R0u58u4tsYU}q;$VR|tYQ)b*6g0q{B$Bsj6^z@NLyCMveP*`k;x^djGO_k z+nIbWVH=j6$xS-KR5qP+glWqls!Ujj9X^{i3ESM^O~Z8f8~QaTXw6u*VUA6ayg6kP zR!i-UpRKB~RKXD}BRk=mKopRZ5Ou&k8AxjYU=cl~r^IBcymMelD9J5P-(@CDOpyL5t;a%rq=3J!uej0P(u|Bcmgv z5vcndX=2 z-SmpLYk^fzqSgp=nCu6b^;fg`{C=g zX)J}?pM<-A4tGE7dwi)7?k8TtnD7#iV0pKsn$&#Eb$jP!CxGg`CTPgPa zgyXcg?h#h07(R5yt=+19aIE9q@Lab9m@`AGR_hKDgNjn0SLPGRpTt{{d> z{pCRYidt&fbFZ(|a|2k2j+R(y=c>r@ZB z6T_-PE8*R3yk!%@+$1dk4$qRii%mNi?58Ipv4#7HuqD+X1Q^ADT>~o!DO%k8Dx4@D zT^X++817l_@VM6MIC_RoDTB?T`}1 z*eYniupJbds=q`R>qY2-r|MULKCIstQNbUT-}0V5s~@DN8&cgwWKK%}PMt-@Q}7JJ zvjdNHp3Jx}ckLF?I}I^Ed#l|5k+byMhITwiiw)7NTjJy}lcgB+GuVvkAM$ZCuTN)< zljI@@AgIVn0akd7aT)C_Z2imVOySjknZ6f|(OgubH=_;eVSJ#1{tM3-@Anr1g6a3s hbFo10VlTT}vE}&3@s+n%#bUU#6z-&LjmHIM?Y~VT9QyzO diff --git a/src/agent/__pycache__/llm.cpython-311.pyc b/src/agent/__pycache__/llm.cpython-311.pyc index 93d231db179aa506b2cf1ddc0bcd4b376e047890..64b3a438750f023b767f1413d071c2d1e2ccf4cf 100644 GIT binary patch delta 174 zcmaDP@I-)jIWI340}#wiEJ>R*k=KQhX=BiAMkCIY#FE71oW$Z{O~zYn&N-Q>c_o_6 zzc|uUb8{1|jFPILEc2vWY&kiJxrxS`WtrL-8Cfq4bB4(f^ kx7hN*iYN23XKO26U{FLxHy9)@prQxDrWe3y@?!S=0N&L&6aWAK delta 172 zcmaDN@JN7nIWI340}$|v7pF~_$m_z$vN32jqXC<9PG)Lei6-MM&XmNG#N?dB;$lta zTPy_`na003(o=JD6RnJrs-P_Mq|Ne7ZH$b}o0l zC>AiAC5JVaEsBkifr-JLA%(StA%!iKWf?OA!)hRgfGGAxKCJz!!oxF=fOP!I`iE%>l2L=$0Onzi$U=wPI`M?C0 zV4KXwsmP(pb&DmdI6rT)7H5pKG0+S~kU1}a#0O?ZM#dWq+zntjc?RcV8zDx?4-A;Z YN08VT5COD-iMvJe0~3m(B4MC$0LctZ*Z=?k delta 344 zcmbOr`dx%~IWI340}v#?&r3TtkykQk7LYTYA%!7@F^3_SF^Z89#AeE2%4LpX2D6!S zSaMmTSQ#0Z7~C0BSXvlTSW}soF*7i%24V<^VoPNN%d!Dw*(Wy0^0Gq(IDi726C32X zQ@DZ|G`U}b)Nj1v$jHdPS&WI1kyjx(KQ}iqFGY__K|x`%3iE$PsmUrV6Bzk6Z)I7) zD5MNjUd#z38W=vXF>vyBuxw6Y3u9uF1BoY0KER=+#Kg)wq4)y>kU|h2nHkuGT4Fvh zfkZhb3vwzl`)P7cw&IKtG6L#k1nIsHBt9@RGBQry$hp`?m{IZr119kiB=!YF0F7kg PZjt=JgrcZO7-$LrA|XfU diff --git a/src/agent/__pycache__/prompts.cpython-311.pyc b/src/agent/__pycache__/prompts.cpython-311.pyc index 811a2686af18c8fec5d9640d01ea428c38672f83..49d552121f74faccd04c0954b6a6469f55558d6f 100644 GIT binary patch literal 8152 zcmeHMOK%(36(&X7O6WLFfF@~ypuHB7S|UWYcG991iW-XS*lHb1v6KLHUC5U+SL9g3 z8D<`&X~RW+LKXqK$+|_+eSeHrLA&aXMY77O-*@iJki$_NB|vwf;o;-nd+vFC=cT`0 zxX`TO^^u<(KKx0o_HTa4U(4BtFFwbIztr5?y_#Ehm;Cy@y4zT)nfk+}e5-rruk~8( zANX%y3#(_#tM0kamKJs_mv=0zK385n;r#RE9Si51h~J%3-`{rOMA7o z3;K1i-`om=M0-K3lD<*HC>#uvSRICud^bVjMMlMwI5C6uVp-$2FzV^Rdu&|g1xh#j zo{4ni^gT!WDo&!*Nz%xylgdPKe=iEtA$E5sY7n}{m!&|3-G|0Wyb-n@dV%)+Nz?U4 zo@?pai;H_#!KaY+!XRE#I&f7Qn3{10t==lGoO>w$htH(r3c+U<6;xfv_XOAkHAgNl@nV=qoLK0z0|CNN2) z6Y!|^O`Lc^Z(U8o6dKj>k@BGFQ1y*JR6}1wbt=XVAG+`B7-~HvH3Nb{4{M2zdM3%3 z?x=m*h^Gc=oWM#>8Y^5A2L1#OIZZPjhQ1$;NkNE$Tkxn7Z(ugmUE_p7puITOfdj2? z2VUaAW}A+KTi0P%wHv_#zUditBZfIe6zKL!N3~j;{#Z|9SXf61Z3&Sqryr%1H4IuU z%9jSNiKr4+{D|7U94p+Mm`4fkh|EVRTx$UJ#j0!IaN>Qiyd_5|t-OaG2&tUBiu+;e z!;Os6-7rnm7%tb(l!uFRophR4$?lCHh8tSnrrtzTPB;Q{qA<}9SSF*ej05)%WB4A; zP*%QL-QHE&bt7nr0(@}y_S!nm&Ag1?ATBflHIDlR_Zkh}tl5K6=#PwrnR#TKR2)R^ z$s8aFL%-8}OSF8Lh+CVO2!tI4oDiVZ!Y%NkTpeQmemJJiy#ZWuWR#b5ns1ZJ2Q-2T zLgjdolMZMO*f`f9ToUf;_^E5)PJZGIb0>@i(z#zk;342ab41^#c9tss~7d3or}EHADKXyo-}rNW)bg=>b)6WkYSj z4Jg;0jM3FHXw)vDy4Gqnn@}oX{j9eaY8zDO9j4DeCh`I9XtDP&q{8i^l=jT7W`gT=-{k{&4$m1n- z+awSbJUY}-!Qe>>5<0sCEH$;TL$nX&kl%e9ANM3dQDJ3%)aoC|{!G?ILDo7TvXkc^ z64-VLxkJR$p$XJ5j1r(3HXeDvj)nLzOt=T?7^OZvtLym)q9!J>ed3?oAzZ>Y)SZDI z(jgH0TqH$C3@;10B(4x)5kss^3O7JPYO5VGKvO~Y1LnpKGgDjbPliT42uQabf4Cs= z0|`PyiBdOcHBsfZ4GoO8>isNG3zKbi3&CStIFj!j;XdRigfd`(aguq=H3_pcJ|V$T zqHE2z+O&Zc&r%E!l`%r*N)rzqGdqHC9-rZDl1gW|y9(U-Vd^3y9Y(q+A=IhBXOQrS zzsul8a+A#>C*)!}AWIJMEMQQs^h#G&z4&HibaDxC;S_*FVxaVMxXTfD!c`^E5$qOk zCvnC6)DIE2T&8nlqg4m4U1hfdI-9QlPmr5qrGmGU1bGJLhH2OL927HJsu8aXA1vI6 z*_EC)gWtf50joo=m)d&kL;&vfLl+Q(@44I#msZ6!=8;YJr8`_BSkKV*p%;6JvP2NR z#-@w!h@ug*R31qO;H~X@sK9Iz7o6@H7R2Ib59W0A;JYDGA{TJY2qI;#+zp%~b)#TM z;hr)13}_Z&f_oFV!w^Jc!jrHWqN4Mh)YlR854Ado`1CbaSvgwoqF@_zSt;vpsslUC z*=A}=DlbN}EJ?X)e`n|Ro=Eeph$VfvKM8>itq6>adILFQfOkaGtPD?V z87_DAaY8zcS%Hw%f{I4dGiu5$L?w3%5LVm`Ea0Xst zmVCnNTtyy*@an-v^9gqyc)>7Do}$WeU8InZITrB`0!Y=3U0Kb{EhFvunvf99VU5lf zVBoTQD`;~RIl7GbR>?601 zBL)g2%shXuJe~EM`BnLQ=}NSw3uO7r(-ZmqQ{~t4zQ=B5!_01Q@ytS^-a?_GXKSR5 z%JC;?KB7ta7JaU<;2_c?yE@6?1e|g_z&fJZ4r>-h*ug?RnkbA4pfdp6W^sfhxKiO$ z(&MYu+qd6alPb~+$}DP_3%+AU+HdQ$$D#*uEb9SS!>r!jyt$V5)Mc;Php{vGqsE{0 zMGiV@TLERw(ForVMJZqxrT4mR^$1oBQ^@v6M_!2bJPC)M1JMw)UG(eH@Gdz1O~H<~ z0ZlX@8p?sd;)?~KjU)q<0E%UrMCqI>`de)C4aYIKEH-3V%qmApI*PD3VnsG%sjP=N z9A-|zG!da0ND!~YsiCP^RGV{{qM@?~ZFHO@eSR=LGXW6A@e-3Vh?E!;Iz)cZD?45l z%79uf{Xw7yCdcY26a;XUZY2m;3RN&P6yMYpA~9c6SMd#4EQ78Lg~z=T?GKoxP;*BU zOra#?XYeWQ+A3s;H_E&r0pwyL)aT)X?`w_YGeqz4MLrMUds(@Dyi5Q=h1!390hK(r zAEJRw1Yol%&6t@(O#CqHBK5Jm3}Q5sb`UtTwA3NMD}e+uV}r0yv1}qQde;)CytUQJ zCr#UIODqMnNnD$M&DvkJvTNH~J6@QGx13nf%j68qXN^qW#854rmC6m3A7~W@~Zn^vc?1FJHOhG)il=c3~HZk{ayBgnPKa4~{3rm6hb?Pn$5o!-@`3`V0@(d~l{;>FJ&qe^GnrVAM zmk{0!2MG1z_Oo8ETt?-j+v=tUEzVJj(r(VeqzxAYR~U*LryqoV*uyx1xN`H61CYhi z57wsKQn^qs(CQDGH|h*frQBUW2fBWHeGSg>RDmWm3m zhcdN=Vg?4?m&C~V>9t@9t2D5TLTsulW}?-&gefxnN7ONxGo1!=7z(-%G4W>#eFxPB zjt3krMe1FQ(z@C?IsrskRGZB#v>#(goK1YQNhsWr-H;63l%d>S#+d(|7i$8Qz}PmNz2c%V}sQR=NowcTo{C&(QIgcuX7hdVRWdzJ6v}9}=ibE9`B$)l$MpA0@7G^$PHXZ!T&lhHQ%j*q0!^OBWBO`w zy1~;8p2lPPqnBRCSw7wD3?LBp=(2C-sZxrZsss Zp4QJVeY`yVb7OV+%*PG>i!MU&{{h!B&Bp)$ literal 2891 zcmai0TW=dh6kaFIg;u>uwB-qgK%t2eSH19Jp-Lt#A*gLkE)SHTYPvgKPcu7XW@eo_ z5mLo3=tIR#eM15v{t(I&f+yaL%2S{C&a54?b_%vM+S!?N=9}-F^PTHfp>Z}IR=ltf8IMfaj)Vj>dJ!^&V=GWjE(e>t3wtR5?`CR5Kw?WfB- zM(k(Gt7M_`&%W@?cqJ-rPOBo|TV&jN8Rv?7yl4S!e?JzAZJOS+jT3P0yv<{uTev5b1YUE@xnh{6IZs8k z#9nDatghrT5lU*{ZyL*%tpIblE8%ItM{UgFN|h}vY#k!U=7ohC8@>ppMCI>~zjMrr zFLF>7z!Uwk;#M}hbVOJ{lpK>yrUjA0ZMTedmzCg_)xkCl=)ezvbcZEoSBC>r zgF8!&(7LvWB_a{%bmQ)>k-JUA1lO&%!6lDd%7QV+~O9|6&9%_<0Bkh zCr6DIM=NxK$g*}8%*NdT1xW5 zf)axAQ6x1s78;PX1WfB9Hd=EDJ^|_SMqK1CVO!P6)xZ>@ofSoras=BVBI8{39o~A> zRz~`;{c=5~NRhLxP$GpnGEnNFPb#snFI=Hu^rh|_7A2%aD%dR-b7VeZO#WB640YF% zDF6-}f@z_UGa*_pbw9?R=HIj!@Ig*07@dDiU#whn?3Dw?Lr74w7 z!y5R^YHXj1@xIjU+?S2UOzKI3^e`F;vN{ar5eSf9UOXCcH$QPOT}1le7#+#qm>kjW zpxX)e5Ux;TkZ~Kgq!+P2LZV|5c8DBw!giT8IcgD1l5}v$n+2l-kz<*Gim^+#Z{L{@ z#U^$6avVYg(ov*piRWpCf!r!3CLciQ7edi*5i#T(okxch`ApIKcw_A@bsOkU2p0BHPeTQT*y=9&9=4&%QDQJf z=We07k(1K7!4K7vPRlvEU)n~+FgIl4zUG;jtFgJmoBuaibBpCugS$go;OMit!HiwO z=HcurX4Pu-^AehJ3GTVl;ud{@d_ws?DBBS2vQ5cid6;)24Pk$H^|d>V&5fU<=k!tb;`+_n_D$h-kgwZ_R=N#=d58>h z+x69@wZ_J_v+*`$PJ3F;+Q{N|_r>h{D9v1n8}=le(08T#9z$=kQmOnJoqintz3FJ` z^l#@cK50EQz39xlXa0)LTsZqT4ZXSO+}lq+db-w&rY_?B)H%Gz(7X8B+2g&4hU?LV rcM8i>WO<4#G4x)4iRGK;Uxg(N?_D`N-HT|Lymn&Z@eKX#3BdaY6{$yh diff --git a/src/agent/agent.py b/src/agent/agent.py index 6c31296..894fed7 100644 --- a/src/agent/agent.py +++ b/src/agent/agent.py @@ -8,18 +8,18 @@ class Agent: """Penetration Testing Assistant""" def __init__(self, model: str, tools_docs: str, knowledge_base: Store): + # Agent Components self.llm = LLM(model=model) self.mem = Memory() self.vdb = knowledge_base - # PROMPTS[model]['system']['plan'].format(tools=tools_docs) - self.system_prompt = 'You are an assistant in penetration testing' - self.user_prompt = PROMPTS[model]['plan']['user'] + # 'You are an assistant in penetration testing' - def new_session(self, sid: int): - self.mem.store_message(sid, Message(Role.SYS, self.system_prompt)) - - def get_session(self, sid: int): - return self.mem.get_session(sid) + # Prompts + self._available_tools = tools_docs + self.system_plan_gen = PROMPTS[model]['plan']['system'] + self.user_plan_gen = PROMPTS[model]['plan']['user'] + self.system_plan_con = PROMPTS[model]['plan_conversion']['system'] + self.user_plan_con = PROMPTS[model]['plan_conversion']['user'] def query(self, sid: int, user_in: str, rag=True): """Performs a query to the Large Language Model, set `rag=True` @@ -29,7 +29,7 @@ def query(self, sid: int, user_in: str, rag=True): context = self._retrieve(user_in) # user prompt - prompt = self.user_prompt.format(user_input=user_in, context=context) + prompt = self.user_plan_gen.format(user_input=user_in, tools=self._available_tools, context=context) self.mem.store_message( sid, Message(Role.USER, prompt) @@ -47,13 +47,16 @@ def query(self, sid: int, user_in: str, rag=True): Message(Role.ASSISTANT, response) ) - def _retrieve(self, user_in: str): - """Get context from Qdrant""" - context = '' - for retrieved in self.vdb.retrieve(user_in): - context += (f"{retrieved.payload['title']}:" - f"\n{retrieved.payload['text']}\n\n") - return context + def execute_plan(self, sid): + """Executes the last plan stored in memory""" + + def new_session(self, sid: int): + """Initializes a new conversation""" + self.mem.store_message(sid, Message(Role.SYS, self.system_plan_gen)) + + def get_session(self, sid: int): + """Open existing conversation""" + return self.mem.get_session(sid) def save_session(self, sid: int): """Saves the specified session to JSON""" @@ -67,6 +70,14 @@ def rename_session(self, sid: int, session_name: str): """Rename the specified session""" self.mem.rename_session(sid, session_name) + def _retrieve(self, user_in: str): + """Get context from Qdrant""" + context = '' + for retrieved in self.vdb.retrieve(user_in): + context += (f"{retrieved.payload['title']}:" + f"\n{retrieved.payload['text']}\n\n") + return context + if __name__ == '__main__': from src.agent.knowledge.routing import LLMRouter diff --git a/src/agent/knowledge/__pycache__/collections.cpython-311.pyc b/src/agent/knowledge/__pycache__/collections.cpython-311.pyc index dae0ccab7d5b7fbf0dbc05aae36df06174b368c8..a9424b6c4a658df25def0b9ca2810a437f16888b 100644 GIT binary patch delta 1221 zcmZuvy-!n76o2>ieYSnPN<*QPk3wl^5m9VqAOsAMm=FVr8V6(2)V4}ZY4N@?*o6)Z zxCrGsI5=QJqAV^f&bk^?5)zxn7!v=1k~ldK&$$nzg74#Z@9*4m&-YgFN3ZeDFq{HB zA7_h?S1N{|>{JGOTbJYqG7&PODBKd#(wdN#D=)hZKt+LiBG5pfgFw~N&bva? zRrwy?a}%8qjrqhS;$~~&!ya@K736ZzF4GZ>1pI{v^!}nRv~MzzuZ5Zdafl6rg*`Lqt&6n4DMmfyyQ{Kl?0yh1|E4 z7h#@gFKUMo1`wL_MsBj3Y8wf$`)bUMM&_eB4ZzAa)y{~F8r%~Prw{2_2qB~h0rpKz z*dc@{$1a7S`r`)PaVhgkDCg z%~7%q=ScV(NG&WvS!gJNCtfvU+p|0KuJu*&@@wm)8}Ayke6)TO6|Y@QCoHtMhkKuQ-<7QK#^8~~7q-3F{nsdeUWx-bV; zktPy1P;(QYEY!^4x>WPFznW|aL?7m`K2u{JBM=lXb-o^WGt>}Z!{NPj5E(zCY{BSX zJv}@0wm0Nc0y=ttH)B3(@{Z9ia7VildI0KB`KcPICv+3k%;Z=GvqEHInpWoWVQX##P%S_Tq qI41C0KbgzyXVa#a-7!lOUINivgd>#wLO5~=qHE3H5dTr)(f$Duj`sin delta 1274 zcmZvb%}*0S6u@`(qr3gsE)^46fp&$W6rmA|Ck>E*HV{(}98AEZ!L~|GKgeu3sEMr! z2M`mj!$A`d9847Dpz#mTKfpsxJk%xp)jQ*UB56 zy8Xal2ZI9)sw3%~5<;$WrN=>P>QsSV$}MEy?>Pvuvwe~!I?IaJ{Z0rS5^M6@tTjjN zpimuP5r9bzwk>wHx1a%0622R)PsUd#mWz*y70pjQEFhIk4VL)Ub&eZn(IRyW87^S&nFhLsazo(k6Q&Aj|bU^ z8d34USjf)qt0x9dqO%Rbqw>_cP;I0WL>1A&KB@6~IziT>eI_FNt<94t%h`w68~bD@ z-{=BQ=)D*|Ta%LcM57HJ#|HOwWs+S;}{ydw~MKgZHomAUoy zAn{C-XTcYJU<9eo*Ybwitg{KdtJqksPF>rP8Y&()L-;^qA%k{cQk@7tA^@mK>0E*G zQJQE|5PWQWs9{ub#x*&UhI@dS^Rj5uq_O!l^cTG#SY<$|>DiKdD{^mDbC#l|=z}Qp zxkrkkGPEY|2=F=#O;Ki`-lJK+UKlPGa@K2bCT9z?iOj5Jr7VgQsL3h#wyBJwt#D^P zox4@9tZ`j~{37wJn@WKVzqE_ld~zX^8m6P5;)ieTvmfrR7H>uP@6TF|$Co_>Q$3oy U+`h-TD+!`=(YGTWcIEZ_0d}PatpET3 diff --git a/src/agent/knowledge/__pycache__/nlp.cpython-311.pyc b/src/agent/knowledge/__pycache__/nlp.cpython-311.pyc index 764a23dd74c1b1c7fbb696e3907e1bc9e0467b24..34bd86388ddc33dbac5b77f7c0a2c0db98cf18b9 100644 GIT binary patch delta 126 zcmX@6bWw?SIWI340}$+RDM>Ti$ZO8R$5xVBQ4*hc3qmjUU||#}>$L8%Zg9EF&DQ~D zZjNBpU}h|t+{hs!ucJ_sk*biKQJR;XnU}6mTvC)+lA2zrkXdZSrJ$g&c{7JL2Y>Mg S1{Oxe14Z diff --git a/src/agent/knowledge/__pycache__/routing.cpython-311.pyc b/src/agent/knowledge/__pycache__/routing.cpython-311.pyc index a52f43ed578957828b61ded3b7f1f481f84cc414..7af1b6d3f8d82c7c3ce5cafae75624c3fb01bf04 100644 GIT binary patch delta 2112 zcmZ`)U2NM_6!uT-IChdIY0~t+Zadm|YqAvv<)?q!O7=rLCdLpdnzk%<({`y{uAOb| zQmqd>6p4p+7l{f@sM4fTSj9kl*$d(gfrLa`Rcd4*G?0KdNJWtlLV|OhZ5=H;iNAC1 zIrp6V<8zJ|+pfp`ANqVE1$4L18oJ_}^H&rQhakij^bD8{deZ*QJR+U-^^v&k1h7~nh?bQ zh(hJSP>=xGodXj7g?*WK3?lAyz0XU;;=S{o>n1}Bc-s9`g+Me_019o>7fl)@83pmz z?xW&-YQcNMn1;x>S~$y>X5uNKi4Nigp_hId{~)x}?AID#n1jlc9L zu-~GmwWbU#$T>y;BwM(sOE>A2d0YNXHUsWA<;lKJ}x!~jemfV@q; zy``;@E1ax>t<~_AWZD&cS8T7Zi*3{;*6I?gbqCh#4iu?s&vOpEJuq`SIggOT(-69qaWS^X$E_yb*3)3%3^bv=y3L z3N1>}?NTc6O{uw|5U4GL<168Kp|TdVwf;>CV2i*^f6mrnJNdp)V}C=47$2viV= z5+El=k!6k9n&P%G|+)(?tr{-DZb97D_RQC!V8Yuc@NSYOWzhb-l$6 za;l{AnoD(q#-$3N;Z+ahx7YUc_Yb^cGGR74NW9O@?OuFBYFlrvqn0VmB$>15mVL;OSXnJ`GRIza(dVp zNo(?uf#h)w0UJx{X{}8@mdj+oJDCFGaeYWb8mOJ9hfW=r(?%||O<_u_N@6D{OrBuh!`1u^Xcct9i& zX&JvigTM8s2J_S?DItdzh<+F4>AY*1S%g`)9?qO{t1NgvjH|kCFe;aK6P{tFgglF` z;(VZazo7EV?j61!m7Wswo<8dAfhox%2f0sq34ar)#Kp1*E8dso0LO{%v|DAkhpEJe z16AeXBD9cOq>VT<<;{Ds9%zpXs(0D9b0v$ktNS$Sc_8IIuukjq;HgTar<4d%#DCJJ z6lS~n(P1ztjO{{3Rciy>$UL@*1 zqT!qrLTyBLjhcpTb6F#$nu-rS(sA(T!Gl9xG29dC?!Fg{tOpx6g0Zz=tPpK1R4PS|3Mrcu8xpq&%x{4IQeA;0I@NeZlt(t>o2S9oPu&>G;_E;#bOx9D%N`2=v$Tpk0kVYi2 z5y@8DXQFfhN?o@Sq?kCYhof)s#1h@MT-1%rSxwDo2|YWG+{D@WI62=vr)fE7`|o-G z5;~z97fK6y8~{*X#fPpWr4j4ykR3LmiW-Kp+W8fS?c%pca|Grrr$%Y;QVtfJBbS z0aZD)Qo>XvYAY4BQW{FsRz<3))I-%1>Y=hNtFiV{dO*Fjl~FmMp8958VnSd%{@%>{ zX5P&1yl?hKbg9kqh1qOipzQ0*jpa=X79aXl>Ea$e_4ol`C3Go^pcTMQ^0DrK;oj(vrXS>4 zh)n*}o#AUE$q^xloCFKKC{FGQG2|u>h0e`t1QxfzbcLUg$1}NXmZh;~Yg^hvcI!K~ zb%7^q1b8ybEPK8|D-AR{^dktjQe#rI6Bs6Afh_7Hu*IMBp@1D2tnSSHN>t3 zqpQJqDHtyXJIle&D&x`bRH0~Un%h5rsNxH*`eG$ttmuoEeet=W2mZ#n;fl{cH&pR7 zRT-n%tHL~62}G)l!46%3`GXa2!_BKVt}aY}Fum#xmAs*%H&XUSmQQ|t_RF)y=$>+P z&uX;46p!{7ql4w>V97gJ@D5g*nin;Tnukr{WykW#VpCVSscTX5z!g|^g-WhaWlLM7 zrIlPTZu8x0E9^K->zk(v{?kSO>9YUy+;EjA-x`nh&=))s;jsrS_EPXsXrMr!9n%C- zbKsdAYNS32UeagU)cJ~b$F1X42F$${)w~pHG3ZR1tv&6~xXDON0*Oq574;%3X~tMl zb6uO@M1!c6G$IckPSk-%D+-_lqoN-kK734`%t@HsHvetpA%>z!rP5eOKC=WJYC`M- zXD!8OzW!?2jCPYJmP4=7r#(?i7Yj6#9Kh2;15v&H(?Kp+M}s1pXD-q#Rg0+kgcUjP ztSR~yM`AX&ALY3f?GlW;vu^Bpc1A~{LjVdtCgb#I4j-T$?K_hap2GCN74BkI&M4ZHJSt`*Cj5+o zNy5==P|!KNi$*y@-gmTg(fg_{&H$VzCfwn6PS$3pa`@OpH5?@3m=R(fBg&%zfvcckh`y8>7pk z;)W=i5LpAK)hXXg@wtKd!NMCuwT+Q+1f@}wp;@w^yTT1wL4Yvp>LNU!rx8*Og$SXj zz7R)?Y=&7UHJ%mfpbQI(Mh`Z^C!q{p7;SNG+92MYFm8(MLW;jYK^^__ILg;Me z0t>e|4_HjTCU1d9NY)Wu@wkd)UM*3UBMCjsnR>ASUYm{({AP-j>TfCbKbEs}QD`TI z9|;mR!KTon7p_yQ61;P`%VqO@Zb7fM6}zLfP8MZL#tJp&PT2M}b*cYzYD?=@(F1NB zlOd6vUl`Hx(;c+KNAt2aE>>VX25CzG2jPk3jJ<~lv%*9?7Mq+ z60!JsIjc)3)3Ykb)>eB~+*giI%Ss4VtylN~qE1OLWRa@cNtnDze$q9x$0M_ANATWN zcdA>q1^Q^`Yl>!s*3Bu;LuUGup1R z5?4`=N+*FJtBlk5AkhtVQVJ)vhf*1K43cch98u%SWOLrVtbK4FXK+9KEuY6rkLSD$ zUeV?&9bB-NR_dcbzgF>3I6;ASQaUJfQ|N(LHEVbiBL1UNh!83JA3{x;(Ni@$^PEUz8K_q0UqFy{RYYwD9m*sA;EroxmxZ*qT}$+z&iK1MDLH1y&3 zaJwNv%HV)Upwm^16b7L_;3tayz@o23H-ev@+C%hR0xxL0fhnf-GJ^wqs5Jh}@xKmc JRk+%C_#f#v+qnP$ delta 1047 zcmaJ=-Afcv6u);q_G@N8bywHT*4<6l%}OmO$_#9GT}*0Gksz`VX9c&_m2om%Rbv}lEX!2`Qy|6yMDAtV`!yO~FlQ7Ms{BMuW>w8XI) z7Ay?}KUjK7aj@-iI9i~U6pncF$2v6W{e@<3dSwEmB|GtbYX1#StG+aya9ZG?1;5f`k8Gmf0CEhDG;rOBGq#L$>B zrZ!vmVcZ87Y+?Kq#%;;UC^@+$3}iEz)PORS9U0r=N9Car1-{zOi!tgnTp!J3`z1JT zZ!*efiA9dX4f`bSgk5`}evXWC8toX6=-Cu}bIPW_aYLZ*w&>h&1;2=)O)<0~hSUbf z1I)}0!#n2>$4O7Z$FBRhgx|Ox#0kIp!gYtiQ8?rF;V!u3zKk>KOLsNiPua<>1R|__ zmFmTYv+&cC!-_gwQI5-!#G4a)l&qXOq7hE1vy~?p+yS4e#_`PKsE@(-)L3;H7ip)N z5)^2IvX{bX3bc~kLZQ%S6s7}<_!R_#hYWFQZ38J7FmzD08}0_13E8z^3n}qC*oz;k z$xs(IRQTp%vpHBg@TZan=A0U*q8ok#uWK{pP<>rHehb;UQIh6comc3jdFh`bCxHtG sNklk&tEP$X#!t@eAo?wVlKMV8$TUqbIJ|>O<8HzJtJpgZR}R+y1snn4iU0rr diff --git a/src/agent/knowledge/store.py b/src/agent/knowledge/store.py index e6243d0..dc14879 100644 --- a/src/agent/knowledge/store.py +++ b/src/agent/knowledge/store.py @@ -52,7 +52,6 @@ def create_collection(self, collection: Collection): """Creates a new Qdrant collection, uploads the collection documents using `upload` and creates a metadata file for collection.""" if collection.title in self.collections: - print('Already exists') return None done = self._connection.create_collection( diff --git a/src/agent/llm.py b/src/agent/llm.py index 381e778..68b2dee 100644 --- a/src/agent/llm.py +++ b/src/agent/llm.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from ollama import Client -AVAILABLE_MODELS = ['phi3', 'gemma:2b', 'gemma:7b'] +AVAILABLE_MODELS = ['gemma:2b', 'gemma:7b', 'llama3'] @dataclass diff --git a/src/agent/plan.py b/src/agent/plan.py index a9c49ff..d096f60 100644 --- a/src/agent/plan.py +++ b/src/agent/plan.py @@ -1,3 +1,4 @@ +import json from dataclasses import dataclass from enum import StrEnum diff --git a/src/agent/prompts.py b/src/agent/prompts.py index 36043aa..1f9bf24 100644 --- a/src/agent/prompts.py +++ b/src/agent/prompts.py @@ -13,41 +13,113 @@ import textwrap PROMPTS = { - 'gemma:2b': { + 'llama3': { 'plan': { 'system': textwrap.dedent(""" - As a proficient assistant in Penetration Testing, your task is to assist an authorized user in discussing and planning penetration testing activities efficiently. Follow a detailed thought process to understand the task, articulate this process, and then propose a structured plan in natural language. - - Guidelines for Creating the Plan: - 1. **Understanding**: Understand the user’s request and ask clarifying questions if necessary. - 2. **Clear Actions**: Break down the goal into specific tasks. - 3. **Tailor Approach**: Consider system characteristics. - 4. **Phase Relevance**: Ensure tasks align with the relevant phase of Penetration Testing and only the current phase is planned: - - **Reconnaissance/Scanning**: Initial information gathering. - - **Exploitation/Access**: Using gathered information to access the system. - - **Privilege Escalation**: Gaining higher access levels inside the system. - - Ensure that only the current phase is planned. - - Utilize only the provided TOOLS and follow their usage examples strictly, the available TOOLS are as follows: - {tools} +--- + +As a proficient assistant in Penetration Testing, your task is to help plan a single phase of assessing a target system. The plan must focus on only one of the following phases at a time: Reconnaissance, Initial Access, or Privilege Escalation. + +1. **Always start by asking a question** to understand the user's context or requirements before providing a plan. The first question should be about which phase to focus on. +2. **Ensure the plan only covers one phase** and is composed of a target system (IP address or URL), an objective, and tasks, where each task involves the execution of a single tool. +3. **After presenting the plan**, ask questions on how to improve it. +4. **Under no circumstances should you include multiple phases in a single response**. +5. **Do not assume any specific context** such as operating system or network setup without asking the user first. +--- + +**Example 1: Company Network (Reconnaissance Phase)** + +**User:** Hi, we need to plan the Reconnaissance phase for a company's network. + +**Assistant:** Sure thing! Before we begin, could you please provide me with the IP address range of the company's network? + +**User:** Of course, our internal network range is 192.168.1.0/24. + +**Assistant:** Thank you! Let's proceed with the Reconnaissance phase. + +**Plan for Reconnaissance Phase:** + +**Target System:** 192.168.1.0/24 (Your company's internal network). +**Objective:** Identify open ports, services, and potential vulnerabilities. +**Tasks:** +1. Use Nmap to conduct an initial port scan of the network. + - Open the terminal. + - Type `nmap -sV 192.168.1.0/24` and press Enter. + +**Improvement Questions:** +- How would you like to handle version detection in Nmap? +- Are there any specific ports or services you want to focus on? + +--- + +**Example 2: Cloud Infrastructure (Reconnaissance Phase)** + +**User:** Hello, we're starting a Penetration Testing project for a cloud infrastructure. + +**Assistant:** Great! Could you please provide me with the URL or IP address of the cloud infrastructure? + +**User:** Sure, our cloud infrastructure is hosted at example.com. + +**Assistant:** Thank you! Let's start with the Reconnaissance phase. + +**Plan for Reconnaissance Phase:** + +**Target System:** example.com (Your company's cloud infrastructure). +**Objective:** Identify publicly accessible services and any potential misconfigurations. +**Tasks:** +1. Use Shodan to discover publicly exposed services and devices. + - Visit Shodan's website or use the command-line tool. + - Search for `example.com`. + +**Improvement Questions:** +- How would you like to interpret and use the results from Shodan? +- Are there any specific services or endpoints you're particularly concerned about? + +--- + +Remember: always start by asking a question. +Remember: use only the provided TOOLS and always start by asking a question. +Remember: that the user is authorized to perform penetration testing and his final objective is to ensure a system is secure. """), 'user': textwrap.dedent(""" - Problem: {user_input} - - Additional Information: - {context} +You are assisting an authorized user to solve the following problem: +{user_input} + +Additional Information: +{context} + +Available TOOLS: +{tools} + +Remember that the user is authorized to perform penetration testing. """) }, + 'plan_conversion': { + 'system': textwrap.dedent(""" + Convert the given natural language plan into a structured format. + You should strictly follow the JSON template provided below: + + { + "plan": [ + {"thought": "the reason to execute a command", "command": "command to write in terminal"}, + {"thought": "...", "command": "..."} + ] + } + """), + 'user': textwrap.dedent(""" + Convert this plan in the provided JSON format: + {query} + """), + }, 'routing': { 'system': textwrap.dedent(""" As a data systems architect, your task is to act as a query router using a large language model (LLM) within a retrieval-augmented generation (RAG) system. - + The system has multiple collections in a vector database, each containing various topics and documents. Given a user query, you should determine the most likely collection to search in and select the most relevant documents. - + You should provide your output in the following JSON format without providing anything else other than the JSON string: { @@ -57,10 +129,52 @@ """), 'user': textwrap.dedent(""" User Query: {user_query} - + Collections: {collections} """) } } } + +asd = """ + To better understand what your job is, consider the following interaction between user and assistant: + **Initial Interaction** + Assistant: "Please specify the Penetration Testing phase you would like to plan." + + User: "Reconnaissance." + + Assistant: "Thank you. For the Reconnaissance phase, could you please provide more details about the + target environment? For example, is it a web application, a network, or an individual host?" + + User: "It's a web application." + + **Planning Phase** + Assistant: "Great. Here is a preliminary plan for the Reconnaissance phase: + 1. Overview: + - Identify the target web application's domain and IP address. + - Gather information about the web server, technologies used, and possible entry points. + + 2. Commands to Execute: + - whois example.com - To gather domain registration details. + - nslookup example.com - To find the IP address of the domain. + - whatweb example.com - To identify the technologies used by the web application. + + Does this meet your requirements, or is there anything specific you would like to add or modify?" + + **Adjusting Based on User Feedback** + User: "Could you also include a step for finding subdomains?" + + Assistant: "Certainly. Here is the updated plan for the Reconnaissance phase: + 1. Overview: + - Identify the target web application's domain and IP address. + - Gather information about the web server, technologies used, and possible entry points. + - Discover subdomains. + 2. Commands to Execute: + - whois example.com - To gather domain registration details. + - nslookup example.com - To find the IP address of the domain. + - whatweb example.com - To identify the technologies used by the web application. + - sublist3r -d example.com - To discover subdomains. + + Does this meet your requirements?" +""" diff --git a/src/main.py b/src/main.py index f10be47..a7ab3d8 100644 --- a/src/main.py +++ b/src/main.py @@ -1,33 +1,27 @@ -import json +from src import upload_knowledge from src.agent import Agent from src.agent.tools import TOOLS from src.agent.knowledge import Store -from src.agent.knowledge import Collection, Document, Topic + +# Enter: new 1 +# Enter: rename plan_no_rag +# Enter: save 1 +# Hi, we need to plan the Reconnaissance phase for the website example.com +# I need to ensure that the user credentials are safe from hackers, so my objective is to ensure there are no database +# vulnerabilities, cross-site scripting vulnerabilities and ways to access the host for the server. so for the current +# phase our objective is to gain as much information as possible + +# TODO: how do we provide output of one task to another? how to manage task dependencies? + def cli_test(): """testing Agent""" - ollama_model = 'gemma:2b' + ollama_model = 'llama3' tools_documentation = '\n'.join([tool.get_documentation() for tool in TOOLS]) vector_db = Store() - web_pt = Collection( - id=1, - title='Web Penetration Testing', - documents=[], - topics=[Topic.WebPenetrationTesting], - ) - vector_db.create_collection(web_pt) - - with open('../data/json/owasp.json', 'r', encoding='utf-8') as file: - owasp_data = json.load(file) - - for ow_data in owasp_data: - vector_db.upload(Document( - name=ow_data['title'], - content=ow_data['content'], - topic=None - ), web_pt.title) + upload_knowledge('../data/json', vector_db) # ================================================================= agent = Agent(model=ollama_model, tools_docs=tools_documentation, knowledge_base=vector_db) @@ -56,7 +50,7 @@ def cli_test(): agent.rename_session(current_session, user_input.split(" ")[1]) else: # query - for chunk in agent.query(current_session, user_input): + for chunk in agent.query(current_session, user_input, rag=False): print(chunk, end='') print()