From 81d8fa07f0d614ad1e9547b4d591d2dedc200cee Mon Sep 17 00:00:00 2001 From: skelly Date: Thu, 19 Sep 2024 22:22:04 +0300 Subject: [PATCH 1/2] Refactored server, fixed double free in client --- client/client | Bin 23656 -> 0 bytes client/src/main.c | 5 ++- server/fraction.py | 8 ++-- server/fractionator.py | 78 +++++++++++++++++------------------ server/main.py | 91 ++++++++++++++++++++++++----------------- server/server.py | 3 +- server/utils.py | 1 + 7 files changed, 101 insertions(+), 85 deletions(-) delete mode 100755 client/client diff --git a/client/client b/client/client deleted file mode 100755 index 8de800a4d0c5fe24b12bd7732ba2db903c6bef18..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 23656 zcmeHv3v^Re`tM20qZCTZ(<&kv1q87r1w<&KDG)d?6biH;3Wl^v+d!MtJfI*|m|Dje zt2m0{_!wq<@ppWnyb5Al5h*G&6y#A9p`+*|78Q}v>M-X1zWqo}V&nSXyYBt3b=UGh z&Ue1&{`R-O{W$xap1VvlrzYrhiaZjO8x@B2HZl@FD@tj;lu#6-GEPasS+5LIx`5mX zKPhJvR}!`)a6Ji=II1K9l3qGpFldZ|morRJB$OT!B)!OC6qg)s>dsLKl{Ai%UT3-# zxRPch47ia2(`5`O5=yync?Ql`3|wBqM^ZRlLaiRfrhOvkx94&_3DZQq#|bF)C|snM zbR*YGn#J`bY!-o#P{xyDLO*G5Tq4(#kYkZ{iq)$WepCuQ38%m&9THL*bjuZ%3S7mgh>vdo%M zW_NfiGb+c8%@{i-)9uV0#atl2C?zPTrp+o)64oeaaLFgxzW7nt zNk8d&YyX;0low}}XFuEj@uY(_^^=1nL-imXlA(j-h06@5YdwA>!{WmWU?M~xf8H;< z5U`&gSAf*}=VZklhdwzD{jNCl{&DEZapSOEe`#KIP?v1=xK0B4_@ps7le-dY>q>(j6+Y2L;r0Y`WtcR$K%kc&vfMH4RPot zap+gYp)ZU>e+cwm%B6~ajuBTFU)dQax1+*kcX*1G3a`hlRM=cDr%Q2rT#llOl}fS8 zW@9Id)e5O%r&40`gs&lG^;VE#QK^L~d7R5g*W)Zx+|Ht9HVFzM=@M=Qoh19(Y7)?QTajO z-plM5I$@QMXG+U#BRG7Fqj6Cltr~hCNd%PbOyjd&Ls#;Q4C{)sFwO8J4P86FrfKLj z_R1q&Lucy>!O&~y+WBgjhE8S6!=Ryua}vyP8oHcANXMw5(- zz?(I6v4mmVB^r89jeMzwPGf{TDl~N2w#2K{(0fT?z%?3rZw>uk4V~sid92gW`$;0; z1`VC&7 z{4a#l6dGK|@IMkxQ)h4u!~Z}yO_{+8hTlgxO_jkV4F4VBG(`sU7`}>dni_*fhPw!- zDKThZ_)@}YDh%owUPw4ifx$F}-$FP|eL;oca|x#@FL>@J5QfYkoTj>9E5j!dPE%a4 znc>-l_a?lN;iCwrrAcrz!>=KnrnulnhF?iIO>Mz-48M$Un$m)67~Y$3n#zI|3{NJU zrm)}=h9?kCQ&%vL;Xh*Mpp>SpppoI{38$$lXkhpkgwqri)B~sX&-!w@@6c_k?-RB9 z%(=XLQ&wZvyJ}O@fIevMep|4CaiKUfb`c`|Qs@ta?`pMHPY@w^L5UWq+ubEhmIw;;g z1*-3}z$!G#zRtITqFem-&&nhDw);p>^;HF$Oc%+;w!W(06iD5!+>q?;RBgJbQ>#q@ zouJ*_=>3P^baC<9q%==mdv-rns>u`}O|{mvM)frXo<_o|wd*TUHYXiYeS7ND3`V7) zA5MXRNN~hq7`oLi1W@(8wd*%fucv;}peKo|5d1pUZQ9T|c9Cj5sWwc)iG?9x_z$f6 zYgi#Da|g=wnOdr=T0&kjY!QYBhYjCMhTnr0H~a?}*64pj=4#hdFY%jN>Z>3OtDDhP z;OSK`m?jKf88*0v46cM;@Oo~Gx*XO2ukgOQQRg!q@tcmQ8`QU26N%iDitw5CtJPKe z74In6+>IMMs(nV-<|eW^g!xiKzL4Hr^uOQ{3>t8kitra4QT@65Ro~x<`w+N3)86W; zy~@fPVY5`&92Be;)lk=z)E|X|Gz;tooYaae}HNz#K0i zU7x8Aj@K#PD^N(HC}e23kajAh41!IjGaxE$rGC?y4y^5g6LqR?4{`^cHcpC;8M4k$JbOMhObz`0*gR)Rq2(_p-ozcP41k>Kt>zMcF$o$T~4K*se zhLK@d+VLYf)e1~$v-nM|@twlRj@e!>evJAzC#k+<)prXn)ij{#s&7Ub@O0I0 z)T_Qp`oM9Fb(jwftdCTiTH#GAdGi=@iAGe^XS%O>=Xf~0D;KBwn{k3uzel>Fdbz0H zj%;4N%i$oeULjRaPmLrR?Aafp8Y~9`wpv?5_NjuSfCjfhTlE(NsQ#+&7}4{4GclP) z-J`WGBDI@LCs|mg`%NdQHTr5=qgnN{@C`%R#Qu3=h^9osGzIJ^tJ-uD)jvtqe^q+^ zMLVknTxtrO!ZqT7tF5rfdU0DSCp<$0>hG$5ar?V$a6vm5XjD>t-d5FDfV9YkOZ3CF zA2bY7WHkxR>oZ-fuDYmrMqp?H={gFR5AJoobv$Bv7s3fiK5Q&x)w$T#nW5pL(3&U; zCjEIm7*NsXz>J+(Hhu~%UP#E>tKl(fpFDCIEf8u;M5$pzX}+q9aQPgx&yjlzDFEb) z_Z-A_Qn*2#2WnWw;Hx?tu4R*{oi#~sWZ)_fU2nGO*3nhSo80g+x%?FQ0wdv)k;_$Nr+t?Jh1OHSesA}rz&r5?0s;s7?j=cZQfsN|+kas}CJ}UZ3puBPF;A0(&Tf8P9 z7i<_pWx)HuIBV)$7~!hdQmV7EjzQQCb5vGXaBx^Kx@}~2{fcj2L&+#Pm3;Ocj|+M= zI{dzdHDS*CVNPB;g#i+X*}evfMZf`O!?TbLOu%9@!EIMs!WXOdr6+!kt&Y1kBD~y>0G!GL`nUojAgBaKEgVF zKb{7$MNLJAxQry>(Vl;y1j?=sN}u}QEIrRxbrN~ojNZ8)y|Zy9O%Z-mbNy?G2^CVy z4GBN`Hf*O@rGr>URS&ME3Rj!jsZQ zx^0aNtfyR6eYKkbX8LM25$3NUjB0n+1%mrjf4`$+DEx0SX@4!#t*vRM@f-6=Xd5X^ z_kA)qB%aAq<-%m`xT1M4&0~5yR2ldEvf1j7q`}q-ss0cVGlh zBZXDpqIy_9go589M-iWMYHc2Q&7Fr*Z7uVRhw~1ItF`Ni;;-R#rw(Nx-x0IE4a#W5 z%my!5%z{soQ8`4-gZThYqV>%zHfTm&^FAuK09^&l;d`iZ#M4F6VmVa`zUQlJtJ;6U z(&^4_`)UrNNTqF3U5J)aBVfqa1PjPdZn}0%0(X!Nb*k@2c=_~s*r7D3t_mor4?3tj zVZ`GF?n0!M`ay7?!ZHu_#B`g?6R(95Pxu9{dP%w{LqTIrF%r;{5!DBNL-s*lD^*)nMqlCG@n+OorBlAs|vw?(bpMuFr} zKrX?Y2VS67U0JQ7XSB6v^Wx>RzlJ3$Ms$P$O(oGZj`j;>5mlSU(P`WC4Jr{y*-ApG z6@H#ZK@7u`Mwjq3x`NuP$#f3w6bklX>se|%8Wx)Z-!L;2Y~JZurt==fGd+b=S!TEJ z%tncMxGmxGG;&8}_McHqI8zt$)g~M7I!IN{9o|u{zDI7MVdjT(mHC0dhk|+V<9{b# zzxtZyv@_LJXOOK+e<53+gYRfpRuUU2KaWn>N7(uX%L0~t5wQl=o>*bYFh=FUS_REAls=m3WK>z&7Emn1Mn&3+wqa?TEZe5-rm$)_bDKNJ++FYmiNW;P zG}x65iEX7}b=K&fE44VXPXRrs=xEHO!VET+i3$0;78!M zT^p^&=05o@Qp@Q6=<qh=UWblsU9)KvcMFmzga|a2zN-VGejy~L6#HAa{Wk{r0RVMt8n}rh++*zI~ezZ z>#MpSt+Yn<@l6Y{pHZCpbz|yq(>nY%7N;IFZH6LbngT9V6^8iC z7jd6>1<#^DzT3d2liGqtTA}z&D2(oL=i=4{DLWYYuS@-Mt^P|=zikDZ0{mDb)fTK1 zYoTH|4)s22EmlD=cMV2IR$AM1wffCU_^5a?(Oc*W>g&)l8&n?>8Mn3m+>JmtsXpA- z`fy+CFW8KJt(T@wV}N5-h-%Fxh&fj08Ipbp~auvdW2fCXVKZDO#+;LDt? zDF(`8sLO0pYuOE4I}KaDgMsRwNeYLcz|7)K3Y#C??0g9OO$QLs)rhFibO3kZK%5NO zk8K7Zg$&v2FE~J@`O^GY0@FoaCx5|Seubc#HsMaADNx253X?l(skE6v-;5l+MqT`9~LtTe%sHtJ-W`vC{tVgu_^VrP=Zh43A@im9LTbkdLhF)6}?zKCD6rR4I zJrbrp!)U`m`%{?qAZQIQh0!%YDQ6G@rR{2@F~x`0i{XLcwXkeH$+Cd~H%qC@4>#-p zD{v|h9u#P79mmJkR@L`5EFt`~CDRlLKnjbBbc$y67g$>q1a`4JB?PNby52=jsNa95 zE^ACvhojvwmNhB}NGm*SVN7N*7lFe4w(dPNV*(s%p$P4VEMQz-QiOL%_0u{Amf^>o zB(jjE_4Nkwm9;u|pKWVksi*sfZY+?u?~{P0a@`5nQU9S$5Skg-4c|4$5bQRRq~Dl6 z-gLyT~Dox+-Il30m@+fc+*#rIB-4OtG0tExC|RSf&otm)mZD6x*vo8Iy~n9{+1$>u zFF&4y#_gxR|Wq zt|cPpvUyx~8wx0PS&BS%r$di7ub06TJ0f!I#YwV=R~?M=Wpm+T5OG#YC^0X6rR?l~M!+UC%;CffX;4GHECm+3d@Sj0UK1 zI`AqeXNb6J6clP5kImu9m}zsAcuKSNvaP3?@==gp&CkypnU$H9k}R9i>B!bEv)L*# zEM@lPwv^-)-T=()SOsP|e`SS@oMtAX${J!#NhUf>WMtf`pO>GLUocNUdG-_&dQm9H~rq zEJqNn`hvMLd8N^M2U z+}?8ZCKsycaYiV1$IodcHiyk+u~G*$jJsK*jzmNp3vDjDg}S!nPemszLMH@{m=?*x z>~2|J2gS+UiiM9Fx~{m^((|ibbGki#>E80L4_;VWaoux8!DrT6pB{DZM?IgdJGt_T z;L+B5PaJ-1#LVs8-dz3Kon15P{`t_rhIQ7o9RttnH!oFh-}&vSocg<~j(m7o+nJBc z=WIK2p{eGB>iHYW`ZR7@W?l6|;eXD)YrAjuB3<7TWnGGjCZ_)Quaubkw;FNe+~-Z1S*r{DgkCgt$5ZOPC5xM|yUWgA|fbnM+XJ~IBW`=w9IT1I&n zoi-gzId*d4#6y35$#7z5wsLCDnmNb!_qQ}RY`d`aH&^Z5GXL@QyAGy4_h#MGbqT|g z?(J-T{X~y_gM-~)`~7O?$X+u`Z~p5*%UcsOinp)JIk~*c?H}oO>kn_ZChh3gzM9u= zAGB?Im*>vlKJrLi;4klXtsL`%lKQ}=ZjYvK==qN~QkLG8F|j=9u`=sJx$}U%Iy~`z&geAdvw?}D%`=lZ|L%+`gsZk3Bxwp7_VZuX)B#JNm_|D-J)^Z|Fyzs{5aG_wHWw>D?2okKDMlJY~}@ zOILM2*z?R&p>7Xf{fd(O+S6UDraYW_?(j27zb`$nyW*!kUF?e&6n_7HiEYj8W6Li6 zXP0G;n=-z}>?Ogk{+;}YlToj?5Wb^QaA>vvpo^nuMw0?!;fcIU%~{=9cj z%ecAcPtX3QW3pzT4oMNe)YQXh^L*GoW@9e&p~6 z_e}ca$T#gmuG{p;7>I&|lE|5&lvxnkOm^v`cey0>C!YUsTQU5XcV*YyfL zZF|uEN@4p4p=B1+!Lt69pCoVm>Tt^O{OV4#_H0QUZy8$h!uS1`H0MsUPyc9z(Yaq6yvuTKkxwizxz|K^+b?b-F7;o&Xwo_}WR z;Lc;4cRk$Y_}_*vIF;W}a^kAYRmEHCW?1$n3@Bav$Oz{(*VlL7`sW9FyglG_=UZyO zPRP&=NSr`_%caMl(0Du)q9+O)PlQ6l0Gk1C0KDhJP^c2{4Zsb6>8C@XM!;&oR=@** zv`h&B>H$}_hC)Wbt$<4ayZj>*S_4=LxDoImU?br8PoNLD1Cai3s^(MZ0~X*W+6cHB zdoW7?PXMj~{1E#~&jRYPJGLJ%6EFZ+1(<}p>`mCKycTc)He+W3mIIanJ_UF`;C0x* zehF|p9wsycUXLd$v~jo`Fb!|R(9xFzMYn3MqN_~P4d~G|X&vfFdp7iXOhrp5M6Vnv zQ@W-R&4Azipb2$lFtP9}J;r0B^h}%DGkr#C_Z3NNl$$P_aLwq!SAxxK(WYD%#P~WQ z%@gXCZA|$H}B(6Xc=HwE}@vo5dR2% zPk=u?ia#aHzX1AP@JDO;4qgQ5_eK4WgTG44Pw05*N5F3cza)xZ80K#UzbkTYX%ycU<{trn4)}LQ@mGfV7r=iN z{J%%>-wpHoB1ab^mzx2@@lAX`fdnW%V?a+q?p}>tjPS2al>B_~`y!_oN7=ub+1GK@F%Sg!K0w9qeyCIVY8G1Q3EHj;p5dS3j zH-bM~%TJgQ_TwV><=~G%Ji_{vZ&&D;4B4Rm>A@D@t6G`F7nuO@soy4|pZ3?5F(={! z_2(bZpQ+!3%UBlSQ$Ozp{S8`v!oy)7H-f(w{I%Nh5v#-%5g+$L<`l++wNW-!hwE?# z{F^XV+@$5Niulk8{p~jJEn0rUA}&C*|AC{mw~>{ySwN~T9cMUDdI+eb8QB%RAWg@( zYm5wD%Fg)v=3I^s5IB2ii0jJ)uFT?mYCAe4|6f0aoP;+f8VT^f&8Ca4Mmj59*jZ5m zf-l$KgM_1a(~li&#}ilDzYQV=nt&aF*+vGDdhDnZ7jhh?7t868e=Bq@$;fa?-z0G9 zx5~K4ApvFiKZSUGN`-t#;4-e55?Fopf-ltn$AJ9%@z?^#iuRo<;Cum#1zaxRJpw)| z;PV1*7jU0|#{~Rb!0!a?(#yz-?l0g_0mlkBRlxZI77Mssz*4yW(V6%QN&P6ps4<4oS!4CX=3)axwRi+0 z&A}K zXJ(FcuXI~&74DJvYx9woLVE^oYi*fD%DlYn(aOBZ*%?{CAT)Vu_PAeQPo4S8>eF(v zvy{B)*<+M>rfh>^pNCVFdv=H2L%uLK1j5P`*Sr9ML}}u$TRrxo43`am@ycss(8to4 z)fghP9z)}$#1PTA)IZbkQ_$h2`3N1!IQ8f894^_7rES#j2st@ z*BHr*4*fBTvOfO>f36q$5_VLUy2~ zpp*Y=t}^0E!4*F~10ebXf<9Z&zX$BdA8Npk==AQh4u6@yZlY0Gh=7jd`!o4|O6mRheG} z!a=X3E~b6k;>iCLhfd%1q4v5~}1pVZ1IGx?|fo@@RCK`3}g07FEaOCbd@{fVuDSY~jYL-1{VCwxg8sLm zMk2HQB7%NqiM4im3h2P!wbG>{S|1-vGEbik7p!)N+3U7hwSE*jtt+)S-I;pb z<|3!nrieziRf?}YMU}{U1O|XYG&?>DWww?f47fqYn+;T3+%TIf>^7FxV2J1Z+@fsG zaPLEB!^;WCB6`!pEM69%{^%BgV_%jkvXm8>Q7jufEG}C(ecU{J$XN0Ht0+fVw3bDY z#ES+|H2%e@D3aNnpF3IfaP|aQQ=+DOMCtJ7>ro^&JhAZyVdbB^f)xHoNogu?Ml&)| zHcH6Pbwx?@uWUt=*b6FAB<^%MJ)OtyxO-(eHCHADguzmIg)zY9s!%f7c1?-Hn_1y< z;!9$ll^Rl^*N%?~+O2|=GiiDTjsBS=UTSfdDw)=m4iv#5z7@tecwNDbr|?mT89bzl zg%re9MVUv*q)}YSWN)8jvXLOu;KqQ|xZBs)J6oVKvRPcP8lC-h7(CSwTefizY!BS8_EO(*~TK8%7 z<$lSNSbx%?XXnT%R(}Cl__I={^uJM*-zfBD`EuVp)_)81yW=PI<-R&S-=rg5K=M(C zpVUt$CJsA(BqH_YK0T`x=yG2km+T-L+T(U0v?nh0(8R>tDUlZ;c7KAmbi`A1}n{%ljR<_`Y5?w2+s2 z0##_=OIyC&?;Ljxmq?L<0>=8kMd-`;%lk!n!?HvMOz7(cl;90qQ2A0{?k~yj$ja}< zO8MCG-;L3)7;j{v72~w6WlD?pYop$sVCc!p23rd)R*5=Y!MqTQjT=#kopq-0Gd`` z?&AgI#*HYSDR-cs0DhEyo{@i>taEVy|^-@v56zJo> zOt&ys>Su;2{DNfZkdXE}scZwIMowdM`1ONQP_QH_;SgNNwjmM2?M5cKUkoTyq#&d6 zu}WMrKx~=&T6)?PjxFiN;S~Uv;V0uKyEU<8?6@R9SLAZI6ZfmyBgP%v`*y59d>+@D K5Tg)-D*p?P-cnou diff --git a/client/src/main.c b/client/src/main.c index 36885a3..a27a753 100644 --- a/client/src/main.c +++ b/client/src/main.c @@ -41,11 +41,11 @@ int main(void) { printf("Connecting to: %s:%s\n", SERVER_IP, SERVER_PORT); sfd = create_sock_and_conn(ainfo); - freeaddrinfo(ainfo); if (sfd == -1) { fprintf(stderr, "Failed to create socket and connect\n"); return EXIT_FAILURE; } + freeaddrinfo(ainfo); if (http_get(sfd, "/", &http_fraction_res) != HTTP_SUCCESS) { fprintf(stderr, "Failed to retrieve fraction links\n"); @@ -85,6 +85,7 @@ int main(void) { } } puts("Downloaded fractions"); + qsort(fractions, lines_read, sizeof(fraction_t), compare_fractions); if (check_fractions(fractions, lines_read)) { // if this works, s0s4 and skelly is to blame! @@ -96,7 +97,7 @@ int main(void) { goto cleanup; } puts("Verified fractions"); - + if (http_post(sfd, "/deadbeef", "plain/text", "{'downloaded':true}", &http_post_res) != HTTP_SUCCESS) { fprintf(stderr, "Failed to send POST request\n"); diff --git a/server/fraction.py b/server/fraction.py index 74ae13e..25e3bd5 100644 --- a/server/fraction.py +++ b/server/fraction.py @@ -4,14 +4,16 @@ from typing import Literal import logging + @dataclass class Fraction: """Dataclass to represent a fraction""" + magic: int index: int iv: bytes data: bytes - + _crc: int = field(init=False, repr=False) _generated_crc: bool = field(init=False, repr=False, default=False) @@ -24,7 +26,7 @@ def header_to_bytes( ) -> bytes: """ Convert the header information of the fraction to bytes. - + endianess: Endianness to use (big, little) include_crc: Include CRC in the returned data (default: True) """ @@ -35,7 +37,7 @@ def header_to_bytes( args = [self.magic, self.index, self.iv] if include_crc: args.append(self.crc) - + header_data = struct.pack(fmt, *args) logging.debug(f"Header data [{self.index}]: {header_data.hex()}") return header_data diff --git a/server/fractionator.py b/server/fractionator.py index c3bafc6..abc9fd9 100644 --- a/server/fractionator.py +++ b/server/fractionator.py @@ -7,20 +7,21 @@ import utils from fraction import Fraction + class Fractionator: MAGIC: int = 0xDEADBEEF CHUNK_SIZE: int = 8192 FRACTION_PATH_LEN = 16 - def __init__( - self, path: str, out_path: str, key: bytes - ) -> None: + def __init__(self, path: str, out_path: str, key: bytes) -> None: """Class to handle loading/preparation of a Fractionator object file to feed to the loader""" - self._path: str = path # Path to LKM object file - self._out_path: str = out_path # Path to store generated fractions - + self.path: str = path # Path to LKM object file + self._out_path: str = out_path # Path to store generated fractions + self._fractions: list[Fraction] = [] # Keep track of the fraction objects - self.fraction_paths: list[str] = [] # Book-keeping of fraction filenames for cleanup + self.fraction_paths: list[str] = ( + [] + ) # Book-keeping of fraction filenames for cleanup # I/O self._buf_reader: Optional[io.BufferedReader] = None # AES-256 related instance attributes @@ -31,29 +32,34 @@ def __init__( self._mode = modes.CFB self._algorithm = algorithms.AES256 self._cipher = None - + @property def cipher(self) -> Cipher: if not self._cipher: - if not self._key or not self._iv: - raise ValueError(f"Missing key or IV (_key:{self._key}, _iv:{self._iv})") - - self._cipher = Cipher(self._algorithm(self._key), self._mode(self._iv)) - + if not self._key: + raise ValueError(f"Missing key (_key:{self._key})") + + self._cipher = Cipher(self._algorithm(self._key), self._mode(self.iv(True))) + return self._cipher + def iv(self, new: bool = False) -> bytes: + if not self._iv or new: + return secrets.token_bytes(16) + return self._iv + def open_reading_stream(self) -> None: """ - Opens a reading stream to the file specified in self._path. + Opens a reading stream to the file specified in self.path. If a stream is already open, this function has no effect """ try: if self._buf_reader is None or self._buf_reader.closed: - self._buf_reader = open(self._path, "rb") - logging.debug(f"Opened reading stream to {self._path}.") + self._buf_reader = open(self.path, "rb") + logging.debug(f"Opened reading stream to {self.path}.") return except FileNotFoundError as err: - logging.error(f"File not found: {self._path}") + logging.error(f"File not found: {self.path}") raise err def _make_fraction(self, index: int) -> None: @@ -65,27 +71,20 @@ def _make_fraction(self, index: int) -> None: data = self._buf_reader.read( Fractionator.CHUNK_SIZE ) # don't use peek, as it does not advance the position in the file - + # generate iv and encrypt the chunk - encrypted_data = self._encrypt_chunk(data) - + encrypted_data = self.do_aes_operation(data, True) + # Create a fraction instance and add it to self._fractions fraction = Fraction( - magic=Fractionator.MAGIC, index=index, iv=self._iv, data=encrypted_data + magic=Fractionator.MAGIC, index=index, iv=self.iv(), data=encrypted_data ) self._fractions.append(fraction) logging.debug(f"Created fraction #{fraction.index}") - def _encrypt_chunk(self, data: bytes) -> bytes: - self._iv = self._generate_iv() - return self.do_aes_operation(data, True) - - def _generate_iv(self) -> bytes: - return secrets.token_bytes(16) - def make_fractions(self) -> None: - """Iterate through the Fractionator object file specified in self._path and generate Fraction objects""" - size = os.path.getsize(self._path) + """Iterate through the Fractionator object file specified in self.path and generate Fraction objects""" + size = os.path.getsize(self.path) num_chunks = (size + Fractionator.CHUNK_SIZE - 1) // Fractionator.CHUNK_SIZE logging.info(f"[info: make_fractions] Creating {num_chunks} fractions.") @@ -110,7 +109,7 @@ def _write_fraction(self, fraction: Fraction): def write_fractions(self) -> None: """Convert the fraction objects to pure bytes and write them in the appropriate directory (self._out)""" - os.makedirs(self._out_path, exist_ok=True) # enusre backup directory exists + os.makedirs(self._out_path, exist_ok=True) # enusre backup directory exists for fraction in self._fractions: self._write_fraction(fraction) @@ -129,10 +128,8 @@ def load_backup(self, backup_path: str): try: with open(backup_path, "r") as f: self.fraction_paths = [line.strip() for line in f] - logging.debug( - f"Loaded {len(self.fraction_paths)} paths from backup." - ) - + logging.debug(f"Loaded {len(self.fraction_paths)} paths from backup.") + except OSError as e: logging.error(f"Failed to load backup: {e}") return [] @@ -150,18 +147,17 @@ def clean_fractions(self) -> None: if not self.fraction_paths: logging.error("No fraction paths detected.") return - + for path in self.fraction_paths: self._clean_fraction(path) self.fraction_paths = [] logging.info("Done.") - def do_aes_operation(self, data: bytes, op: bool) -> bytes: """Perform an AES-256 operation on given data (encryption [op=True]/decryption [op=False])""" - if not self._key or not self._iv: - raise ValueError(f"Missing key or IV (_key:{self._key}, _iv:{self._iv})") + if not self._key or not self.iv: + raise ValueError(f"Missing key or IV (_key:{self._key}, iv:{self.iv})") cipher = self.cipher operator = cipher.encryptor() if op else cipher.decryptor() @@ -169,11 +165,11 @@ def do_aes_operation(self, data: bytes, op: bool) -> bytes: return operator.update(data) + operator.finalize() def close_stream(self) -> None: - """Closes the open stream to self._path and resets self._buf_rw_stream""" + """Closes the open stream to self.path and resets self._buf_rw_stream""" if isinstance(self._buf_reader, io.BufferedReader): self._buf_reader.close() self._buf_reader = None - logging.debug(f"Closed stream to {self._path}.") + logging.debug(f"Closed stream to {self.path}.") return logging.debug(f"No stream was open.") diff --git a/server/main.py b/server/main.py index 0b770fb..cf480cb 100644 --- a/server/main.py +++ b/server/main.py @@ -22,9 +22,7 @@ def handle_args(parser: argparse.ArgumentParser): """Configure the given ArgumentParser""" - parser.add_argument( - "--file", type=str, help="Path to LKM object file to use" - ) + parser.add_argument("--file", type=str, help="Path to LKM object file to use") parser.add_argument( "-b", "--bind", @@ -50,22 +48,43 @@ def handle_args(parser: argparse.ArgumentParser): parser.add_argument( "--rm-backup", action="store_true", help="Remove the generated backup file" ) + return parser.parse_args() + + +def validate_lkm_object_file(file_path: str) -> str: + """Validate file type and existence.""" + if not os.path.exists(file_path): + raise FileNotFoundError(f"File '{file_path}' does not exist.") + + _, ext = os.path.splitext(file_path) + if ext != ".ko": + raise ValueError(f"Invalid file type: {ext}. Expected a '.ko' file.") + + return os.path.abspath(file_path) + + +def generate_aes_key() -> bytes: + """Generate a 256-bit AES key.""" + key = secrets.token_bytes(32) + logging.debug("Generated AES-256 key.") + return key + + +def handle_cleanup(fractionator: Fractionator, backup_path: str) -> None: + """Clean up fractions and remove backup file if necessary.""" + if os.path.exists(backup_path): + fractionator.load_backup(backup_path) + fractionator.clean_fractions() + try: + os.remove(backup_path) + logging.info(f"Backup file '{backup_path}' removed.") + except FileNotFoundError: + logging.critical(f"Backup file '{backup_path}' not found.") + else: + logging.warning(f"No file found at '{backup_path}'.") -def handle_cleanup(fractionator: Fractionator, path: str) -> None: - fractionator.load_backup(backup_path) - fractionator.clean_fractions() - try: - os.remove(path) - except FileNotFoundError: - logging.critical(f"{path} is not a valid file.") - return - if __name__ == "__main__": - parser = argparse.ArgumentParser() - handle_args(parser) - args = parser.parse_args() - # ensure dual-stack is not disabled; ref #38907 class DualStackServer(ThreadingHTTPServer): def server_bind(self): @@ -75,36 +94,32 @@ def server_bind(self): return super().server_bind() def finish_request(self, request, client_address): - self.RequestHandlerClass(request, client_address, self, - directory=args.output) + self.RequestHandlerClass( + request, client_address, self, directory=args.output + ) - key = secrets.token_bytes(32) - logging.debug(f"Generated AES-256 key.") + parser = argparse.ArgumentParser( + description="Erebos Server: Prepares and stages the LKM over HTTP" + ) + args = handle_args(parser) out_path = os.path.abspath(args.output) - backup_path = os.path.join(out_path, BACKUP_FILENAME) # backup file path - - - fractionator = Fractionator("", out_path, key) - + backup_path = os.path.join(out_path, BACKUP_FILENAME) + + fractionator = Fractionator("", out_path, generate_aes_key()) + handle_cleanup(fractionator, backup_path) - if args.clean: exit(0) + if args.clean: + sys.exit(0) - if not args.file: - raise ValueError("The --file flag is required for this mode.") - file_path = os.path.abspath(args.file) - _, ext = os.path.splitext(file_path) - if ext != ".ko": - raise ValueError(f"Invalid file type") + file_path = validate_lkm_object_file(args.file) - # TODO: Implement path validation - fractionator._path = args.file + # Set up Fractionator with the provided file path + fractionator.path = file_path + # Prepare the fractions fractionator.make_fractions() fractionator.write_fractions() fractionator.save_backup(backup_path) - - - # lkm.close_stream() - # Stage fractions over HTTP + # Start the server for staging fractions start_server(DualStackServer, args.port, args.bind) diff --git a/server/server.py b/server/server.py index e226e7a..fbdd037 100644 --- a/server/server.py +++ b/server/server.py @@ -7,6 +7,7 @@ import io import os + class PlainListingHTTPRequestHandler(SimpleHTTPRequestHandler): """Lists the links to the files in the given directory in plain text""" @@ -52,7 +53,7 @@ def list_directory(self, path): return f -def start_server(ServerClass, port: int=8000, bind=None): +def start_server(ServerClass, port: int = 8000, bind=None): test( HandlerClass=PlainListingHTTPRequestHandler, ServerClass=ServerClass, diff --git a/server/utils.py b/server/utils.py index f2ce9c5..7bd6041 100644 --- a/server/utils.py +++ b/server/utils.py @@ -1,6 +1,7 @@ import random import string + def random_string(n: int = 16, sample: str = string.ascii_lowercase + string.digits): """Returns a random string using the characters defined in sample""" return "".join(random.choices(sample, k=n)) From 0e8be5a59999d69ac1c53d409b8fc82817cc4d06 Mon Sep 17 00:00:00 2001 From: skelly Date: Sat, 21 Sep 2024 22:35:02 +0300 Subject: [PATCH 2/2] Added AES_CFB_HELPER class to seperate AES functionality from Fractionator class, refactored Fractionator. --- server/fractionator.py | 185 ++++++++++++++--------------------------- server/main.py | 2 +- server/utils.py | 30 +++++++ 3 files changed, 93 insertions(+), 124 deletions(-) diff --git a/server/fractionator.py b/server/fractionator.py index abc9fd9..1985c9c 100644 --- a/server/fractionator.py +++ b/server/fractionator.py @@ -1,141 +1,106 @@ -import os import io -from typing import Optional import logging -import secrets -from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes -import utils +import os +from typing import Optional + +from cryptography.hazmat.primitives.ciphers import algorithms + from fraction import Fraction +import utils -class Fractionator: +class Fractionator(utils.AES_CFB_HELPER): MAGIC: int = 0xDEADBEEF CHUNK_SIZE: int = 8192 - FRACTION_PATH_LEN = 16 - - def __init__(self, path: str, out_path: str, key: bytes) -> None: - """Class to handle loading/preparation of a Fractionator object file to feed to the loader""" - self.path: str = path # Path to LKM object file - self._out_path: str = out_path # Path to store generated fractions - - self._fractions: list[Fraction] = [] # Keep track of the fraction objects - self.fraction_paths: list[str] = ( - [] - ) # Book-keeping of fraction filenames for cleanup - # I/O + FRACTION_PATH_LEN: int = 16 + algorithm = algorithms.AES256 + + def __init__(self, file_path: str, out_path: str, key: bytes) -> None: + """Prepare a Fractionator object for reading and generating fractions.""" + self.file_path: str = file_path + self.file_size: int = 0 + self.out_path: str = out_path + + self._fractions: list[Fraction] = [] + self.fraction_paths: list[str] = [] + self._buf_reader: Optional[io.BufferedReader] = None - # AES-256 related instance attributes - self._iv: Optional[bytes] = None # AES-256 initialization vector - self._key: Optional[bytes] = Fractionator.validate_aes_key( - key - ) # AES-256 cryptographic key - self._mode = modes.CFB - self._algorithm = algorithms.AES256 - self._cipher = None - - @property - def cipher(self) -> Cipher: - if not self._cipher: - if not self._key: - raise ValueError(f"Missing key (_key:{self._key})") - - self._cipher = Cipher(self._algorithm(self._key), self._mode(self.iv(True))) - - return self._cipher - - def iv(self, new: bool = False) -> bytes: - if not self._iv or new: - return secrets.token_bytes(16) - return self._iv + + super().__init__(key, self.algorithm) def open_reading_stream(self) -> None: - """ - Opens a reading stream to the file specified in self.path. - If a stream is already open, this function has no effect - """ - try: - if self._buf_reader is None or self._buf_reader.closed: - self._buf_reader = open(self.path, "rb") - logging.debug(f"Opened reading stream to {self.path}.") - return - except FileNotFoundError as err: - logging.error(f"File not found: {self.path}") - raise err + """Open a stream for reading the object file.""" + if not self._buf_reader or self._buf_reader.closed: + try: + self._buf_reader = open(self.file_path, "rb") + logging.debug(f"Opened stream to {self.file_path}.") + except FileNotFoundError as err: + logging.error(f"File not found: {self.file_path}") + raise err def _make_fraction(self, index: int) -> None: - """Read from the object-file and generate a fraction""" - if not isinstance(index, int): - raise ValueError(f"index must be an integer (got `{type(index)}`)") - # Open a stream to the file and read a chunk + """Generate and encrypt a fraction from the object file.""" self.open_reading_stream() - data = self._buf_reader.read( - Fractionator.CHUNK_SIZE - ) # don't use peek, as it does not advance the position in the file - # generate iv and encrypt the chunk - encrypted_data = self.do_aes_operation(data, True) + data = self._buf_reader.read(self.CHUNK_SIZE) + encrypted_data = self.encrypt(data) - # Create a fraction instance and add it to self._fractions fraction = Fraction( - magic=Fractionator.MAGIC, index=index, iv=self.iv(), data=encrypted_data + magic=self.MAGIC, index=index, iv=self.get_iv(), data=encrypted_data ) self._fractions.append(fraction) + logging.debug(f"Created fraction #{fraction.index}") def make_fractions(self) -> None: - """Iterate through the Fractionator object file specified in self.path and generate Fraction objects""" - size = os.path.getsize(self.path) - num_chunks = (size + Fractionator.CHUNK_SIZE - 1) // Fractionator.CHUNK_SIZE + """Generate all fractions from the object file.""" + self.file_size = ( + os.path.getsize(self.file_path) if not self.file_size else self.file_size + ) - logging.info(f"[info: make_fractions] Creating {num_chunks} fractions.") + num_chunks = (self.file_size + self.CHUNK_SIZE - 1) // self.CHUNK_SIZE + logging.info(f"Creating {num_chunks} fractions.") for i in range(num_chunks): self._make_fraction(i) - def _write_fraction(self, fraction: Fraction): - """Write a fraction to a file""" + def _write_fraction(self, fraction: Fraction) -> None: + """Write a fraction to a file.""" path = os.path.join( - self._out_path, utils.random_string(Fractionator.FRACTION_PATH_LEN) + self.out_path, utils.random_string(Fractionator.FRACTION_PATH_LEN) ) - with open(path, "wb") as f: - header_data = fraction.header_to_bytes() - data = fraction.data - - f.write(header_data) - f.write(data) - + f.write(fraction.header_to_bytes()) + f.write(fraction.data) self.fraction_paths.append(path) logging.debug(f"Wrote fraction #{fraction.index} to {path}") def write_fractions(self) -> None: - """Convert the fraction objects to pure bytes and write them in the appropriate directory (self._out)""" - os.makedirs(self._out_path, exist_ok=True) # enusre backup directory exists + """Write all fractions to disk.""" + os.makedirs(self.out_path, exist_ok=True) for fraction in self._fractions: self._write_fraction(fraction) - def save_backup(self, backup_path: str): + def save_backup(self, backup_path: str) -> None: """Save fraction paths to a backup file.""" try: with open(backup_path, "a") as f: - for path in self.fraction_paths: - f.write(path + "\n") # Ensure each path is written on a new line + f.writelines(f"{path}\n" for path in self.fraction_paths) logging.debug(f"Backup saved at {backup_path}.") except OSError as e: logging.error(f"Failed to save backup: {e}") - def load_backup(self, backup_path: str): - """Load fraction paths from the backup file and initialize self.fraction_paths.""" + def load_backup(self, backup_path: str) -> None: + """Load fraction paths from a backup file.""" try: with open(backup_path, "r") as f: self.fraction_paths = [line.strip() for line in f] logging.debug(f"Loaded {len(self.fraction_paths)} paths from backup.") - except OSError as e: logging.error(f"Failed to load backup: {e}") - return [] + return - def _clean_fraction(self, path: str): - """Delete a fraction file""" + def _clean_fraction(self, path: str) -> None: + """Delete a fraction file.""" try: os.remove(path) logging.debug(f"Removed {path}.") @@ -143,45 +108,19 @@ def _clean_fraction(self, path: str): logging.debug(f"File not found: {path}") def clean_fractions(self) -> None: - logging.info("Cleaning fractions . . .") - if not self.fraction_paths: - logging.error("No fraction paths detected.") - return - + """Delete all written fractions.""" + logging.info("Cleaning fractions...") for path in self.fraction_paths: self._clean_fraction(path) - - self.fraction_paths = [] - logging.info("Done.") - - def do_aes_operation(self, data: bytes, op: bool) -> bytes: - """Perform an AES-256 operation on given data (encryption [op=True]/decryption [op=False])""" - if not self._key or not self.iv: - raise ValueError(f"Missing key or IV (_key:{self._key}, iv:{self.iv})") - - cipher = self.cipher - operator = cipher.encryptor() if op else cipher.decryptor() - - return operator.update(data) + operator.finalize() + self.fraction_paths.clear() + logging.info("Cleaning complete.") def close_stream(self) -> None: - """Closes the open stream to self.path and resets self._buf_rw_stream""" - if isinstance(self._buf_reader, io.BufferedReader): + """Close the file stream if open.""" + if self._buf_reader: self._buf_reader.close() self._buf_reader = None - logging.debug(f"Closed stream to {self.path}.") - return - - logging.debug(f"No stream was open.") - - @staticmethod - def validate_aes_key(key: bytes) -> bytes: - """Check if key is a valid AES-256 key (32 bytes)""" - if not isinstance(key, bytes) or len(key) != 32: - raise ValueError( - f"Invalid AES-256 key. (expected 32 bytes of `{bytes}`, got {len(key)} of `{type(key)}`)" - ) - return key + logging.debug(f"Closed stream to {self.file_path}.") - def __del__(self): + def __del__(self) -> None: self.close_stream() diff --git a/server/main.py b/server/main.py index cf480cb..4148459 100644 --- a/server/main.py +++ b/server/main.py @@ -115,7 +115,7 @@ def finish_request(self, request, client_address): file_path = validate_lkm_object_file(args.file) # Set up Fractionator with the provided file path - fractionator.path = file_path + fractionator.file_path = file_path # Prepare the fractions fractionator.make_fractions() fractionator.write_fractions() diff --git a/server/utils.py b/server/utils.py index 7bd6041..58654d5 100644 --- a/server/utils.py +++ b/server/utils.py @@ -1,5 +1,35 @@ import random import string +import secrets +from typing import Optional + +from cryptography.hazmat.primitives.ciphers import Cipher, modes + + +class AES_CFB_HELPER: + LENGTH_IV: int = 16 + + def __init__(self, key: bytes, algorithm) -> None: + self.algorithm = algorithm + self.mode = modes.CFB + self.key: bytes = key + self._iv: Optional[bytes] = None + + def get_cipher(self, iv: bytes) -> Cipher: + """Return a cipher instance.""" + return Cipher(self.algorithm(self.key), self.mode(iv)) + + def get_iv(self, new: bool = False) -> bytes: + """Generate or reuse initialization vector (IV).""" + if not self._iv or new: + return secrets.token_bytes(self.LENGTH_IV) + return self._iv + + def encrypt(self, data: bytes) -> bytes: + """Encrypt data using AES-CFB mode.""" + cipher = self.get_cipher(self.get_iv(True)) + operator = cipher.encryptor() + return operator.update(data) + operator.finalize() def random_string(n: int = 16, sample: str = string.ascii_lowercase + string.digits):