From 5510d22cbd212d8cd93b7aa8e7573ced9a52986b Mon Sep 17 00:00:00 2001 From: Zarithas Date: Sun, 9 Nov 2025 21:06:07 -0500 Subject: [PATCH 1/4] Checking WIP Agent Movement Workflow --- AirlockTools_Server.py | 93 ---- IRT_icon_32-512.ico | Bin 21181 -> 0 bytes Server/scheduler_async.py | 215 --------- README.md => docs/README.md | 0 flows/localApproval.py | 455 ++++++++----------- screens/moveagentworkflowscreen.py | 61 +++ screens/policyselectorscreen.py | 91 ++++ utils/selector.py | 6 +- utils/tui.py | 60 +++ utils/utils.py | 38 -- widgets/OTP_generate.py | 13 +- widgets/agentmoveoperations.py | 704 +++++++++++++++++++++++++++++ widgets/multiagentselector.py | 12 +- widgets/policyselector.py | 505 +++++++++++++++++++++ widgets/resultsdisplay.py | 178 ++++++++ widgets/retro_terminal_theme.py | 38 ++ widgets/themeselector.py | 25 +- 17 files changed, 1868 insertions(+), 626 deletions(-) delete mode 100644 AirlockTools_Server.py delete mode 100644 IRT_icon_32-512.ico delete mode 100644 Server/scheduler_async.py rename README.md => docs/README.md (100%) create mode 100644 screens/moveagentworkflowscreen.py create mode 100644 screens/policyselectorscreen.py create mode 100644 widgets/agentmoveoperations.py create mode 100644 widgets/policyselector.py create mode 100644 widgets/resultsdisplay.py create mode 100644 widgets/retro_terminal_theme.py diff --git a/AirlockTools_Server.py b/AirlockTools_Server.py deleted file mode 100644 index 183635c..0000000 --- a/AirlockTools_Server.py +++ /dev/null @@ -1,93 +0,0 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . - - -# TODO Add CSV injection prevention -# TODO Continue OTP and Local approval rewrites -# TODO Explore pywin32 -# TODO Fix Requirements.txt -# TODO Create Generic system_config.json for gitea - - -import logging -import os - -import dotenv -import urllib3 - -import flows.localApproval as la -from Server.scheduler_async import ( - recurring_job, - register_function, - reload_jobs, - start_scheduler, -) -from services.API import AirlockAPIWrapper -from services.policyhandler import updateAuditPoliciesFromEnforcementPolices -from services.security import getAPI -from utils.setup import setup - -urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) - - -def main(): - - # Determine working directory, setup directory, configure logging, sent env, get API and URL if not already stored - - working_dir = setup() - - logger = logging.getLogger(__name__) - - dotenv.load_dotenv(dotenv_path=working_dir / ".env") - - try: - url = os.getenv("URL") - username = os.getenv("USERNAME") - - if not url: - raise ValueError("Missing URL in environment variables.") - if not username: - raise ValueError("Missing USERNAME in environment variables.") - - logger.debug(f"Retrieved URL: {url}") - logger.debug(f"Retrieved Username: {username}") - - except ValueError as e: - logger.error(f"Configuration error: {e}", exc_info=True) - raise - - api = AirlockAPIWrapper( - base_url=str(os.getenv("URL")), - api_key=getAPI(username, "AirlockTools"), - ) - - logger.info("Running non-interactively to start monitoring Airlock Changes") - - register_function("monitorLA", la.scheduleAddingLAHashes) - register_function("updateAuditPolicies", updateAuditPoliciesFromEnforcementPolices) - - if not os.path.exists("scheduling\\jobs.json"): - recurring_job("monitorLA", "monitorLA", interval=50, unit="seconds", args=[api]) - recurring_job( - "updateAuditPolicies", "updateAudit", interval=5, unit="minutes", args=[api] - ) - else: - reload_jobs() - - start_scheduler() - - -if __name__ == "__main__": - main() diff --git a/IRT_icon_32-512.ico b/IRT_icon_32-512.ico deleted file mode 100644 index 103b72ef73c7c58a81ad42b075aad1a626c2402f..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 21181 zcmafZWl&r})9$jn_~LHC-8HzoTL=~$0zrd2i@OGbyC!Jxpo;`Zf)g~z;tq>D+cS08pp@(o6UN00H5@{lb6gNp=9>m8a()^ngX#ey9 zsD=}B007XeioA@zZ(hsi3CDK|gu_z&9;FW9T(b7MB|(Y2$MMA&A52Kz;32cd(6!@| z=;{w@n4+n^BP!E()0=v32$4}^iPH_?Af$^Rt^pUTCN`%}<_s>SnU}S2pRR6NjWb-- z-Sm(k<4~UDXY=nqPa24;em!Y*J@z)BVI zZ|H7LO~}HKQ4p{U_KCL@X^GI3E+4rE92HtjT%-Uq@#1|=Unuxw23U%vd#J(YsGjq% z(Y2tw1uS6IxXRjKx`7SjcIHDl`f=eIzSATha6!@6e2F%}C@0$2inGAKh+;0#_)Y%Y zcPAh*y*BHldh8>zcvacEOE1`akqz6PuvOVMZ!N+dzD*vXx}vi$X!5VLE5{xOb#N^f zfXQD$#qu=@dC`wE#wmCSG4nwWxG;S-by2l3qZ++{zxN0iRN>AjG@GyUXDV~blK@1* zQVk=pZzloC4(f8u^Uc~?8vOM}`Yp4&KPCUdgwq;Rut}XIvarvqiu%4J=jB~|(Ha?f z9xDq=q~Pz`?dMQ0Jd5IlFCNbiceL_c^oR4QL09N>{4u7oHsskN;Ya?Y+2jgaudYXun5)z8$9v4K zH<@w2igsUMmqPVWH_4bDMtrvpC9&_a|HV}le``%$cahQ-W`BPFw9B zLeSs0h6Sf>;a0zw-oS%*;w$~WMa)C1r7OqDM(`#+b)$>t53r{Dx{*klOz0KlF(dFO ze=SDrq4Qvg*GX<>x?9UJ{)kqB^3H-^Mkv8TAPMPcAUB-4MPjK=sfPeg0mrh?{c9Ph zv7xF={L_p)d7TJ;l~@15vA#}*_G76+NQwVcxj z=jeRaDh&CbodZ-9H0A4LEyDgE-hk(ynB_m-pgr2}BLIM6{aI}gjo;_QF^|~kMB<^Q+3)R%Ufgu#l9@(YE@Qo=tO!4=^Hku<1_!;;ah=U|&P1A*+5DV%> zHbgNU2bzLdA0E25Pd{f!`4|-y+a}a#^Vub!j04vI6M)TlKt9&3uZ4(PISJf_fWD$* zL1A$?$<36fF^Z82RPF`r1-b%ntnEtMC_jeI1K2}91EHNL@-U%cGC;!MVApslOps;U zSAMCVLwZvmRMS7`Q4DKJ4C<`DWF798O>fsgc@anTMNYSrE$^0x*Cn73cY>VUIM@}9 zY)nv@&;x3)lXPD3to`!J0Eic$b}HPt!Mi|JgtncOS~rSmB%RicKZN=gfC+jq2Ds>7 zqa&2^<5N6L$xNi4pk8Js6AnKEZ(MB(z|9f~-uF-3Bv{Sov6SLS`%mTd42F)N(R*_N zWTvlAEV4>?GegOio^t#eWC-y|@)BHuG9b11k04RtM$>j^5=}x}m<~b8YD^qPq5c9V z0n(e{@g}0C%;p=w`#W9Q8T9Og?8qKU8So+2D>w4?RG_I+g`e={qO8=8Pn|LFq_FRAyF(~602P8(`geyu%5AR3T= z1@J4ph|*YWJ#t=SZGm^Pfk;;+i|3Wu0}W@$B`T>gRzVCMIdwUUYe!~ZYl^yJb z(Pwt$0#N7A>Av#Bb}1qKRpC}c)+N=*Vc11#mYR`{DyLulSc&@j2;7c%1k8AX5;b=2CebN&VI?=Q*z@b}&@E5Y>1z9$ zKE7y$uS0nmi!eO0u~`wywS-GIQw5E=d9n|`E;|DYlyY_V;?76MbhkDi z{c=oMap37bHvA?|*^Mi)Y9dWOtUsDyh87qIitJ8%R*e(f3(g&BJ82{$>7ayfH`b9J z&2C+_-OWc#Iw#p%RFPlor1u_>= z5UL54gxwxBbv3ghP%>-})_KX)#S>vA{8PK>%p9w_`mo*VRT!p<U@|+E?cOqr6Qe_G~JmcnthJDX}bPv^}wD7*)Vjyz@ zRLCylftY_A^@JX(%&Y(4H5bEX9YlAAh4Y#sr#kLa!J=W23pA<7Y0V&2=QsC)JQD{j z-2DJQA3R^!uMAJ0vytWw3}1Y`1&K@$^lwfmZ<3tQbO74-A1&coZ+Mw+WhbH=<03TI zv`2X#UDWd%$@R?1ql|Rk1(ikQ1D^1G(fL8aPRo+1fv6)y@#yL& zkiI|hWn36`Ol9PpwkkeM<0;r8dk@2(@)*{MQeHNrx&pF16FNmE(Q9B4Rfre!wJraY z%yJd`F14M5S(HH89p)O{>hLE;D|yEKK`BUA=-n#gA9HG)2X%bOPrPb-9O_AEcbBwt zFQ6svBJ&D6?bH7yHc9`J3;icH`QR4<0RXh6|Hh`3T*riWZ-|CNc0ay%5#!8n)umH) z{?;;EL8Q!n&PF&TZ%j+79j;MoA#fWLhO3j4QU%^D{SdQ{s%OrU7#lamLaySRVQjCa zR?IdT>TO&+@vSwV&x3l`NvfuM=O|!je0f-8K8$@)#Akdlp!;8-dZHvP1k!*XQLfCe zwrD%pu7Ecn^KT*1`zV{d7*JMgG{o`gTg)##s@~Yu?WK*^`C+sc5(pjm zqo|jh$&0#3saf(G`sF85GbFkqh^ftCVyxZ7X9pNwbS#n?(lc-w#9=9r_ncXrCKQN) z@Ux(f?Ri7v9ZUgtC>jw)GpML`4sUd>fh?bE!2@6z7WbrG=ZZ=a1E+v%T}NMirM0-A zRqMY=7|6#@X5jmeaXL=fZ%M}<(GWZc6?m-|APX{)ph)z;OGl2?0W~27)DlM;Jl{E>I zUND@hNYy!f6Eh-sw}E2m#|0(DVl5E&mVNtgW|Zfh5fDR<=UzNNhDyML47+t#b{1VJ z+G7tWsx!YNCU29?jdO9`T+0wWYk04{PHBn46@;xP_IcNcaGlyG(rl9@A5|e{Dj77f zxTX^>#2c0tLgSH86f0!3_d(fCA#>U|?1TJ}WLwL~{K9>2Q!mB%%3qq~AbVZ&s8tN>BhR%1S50vj@5lz-M6{wG7?HFgOs|_;^OJH%eW78t0 z9toU;K0qdo5}DBUTnZDNgXYqi`^gn0t_;8_Z3v z=|5`N6kXjAfFs^2xPG&KXfd=nSP^%cXc{_@zReBH0KGApsVu2hxENW{_cdxB&O-8`j^Q1)k$@a!lg zES{zZCoIPE2Jq|%?@;&{QT2bmnnPc!Sf(*&ler>IM=cTnyb+7^vqk%=gdUZQuQ-hM z1y`4mfo@B{@``jl<2PE!y`ZdkzwlRuViD|Fr1i1=u4by@s{Q^-i^Y(cj#Ce-H`Z^8zTPu;Dbr`(FltugQ zNhWteT7b%tj&t(*a=gtXoc?n5^|&`Y!W{VVpWs1nTR2;&rqk~jc#h~r;40@UxMZ~R z+)SU>PqzmUV_vwIvYz-o2=pv|wGFB0!8r~I4D_QA1w}CFAb|5ET9~-q5&uHm@Pb}= z5;a15FX};x-H$FYIJ5F+0VqMa4DyJMkNSw8{T{h_KY|Awf*&EFot1uR^Hyf{K_cxg zF_q7k+&~J7Zl-kfCf4;4z+sM}V=8uerQUqZ{Iek|)w83QOyoVntBFq*o5z4wq~os1 z)944UwJ1j(&PSWqAfKNCzB}J}ZUTAm(%Bnqca=-{3!>XRzX|`wTu0!Ip6s5bSMK(S zaYNA+5I|uBzIgBJ)J2qiI(wo<6A5=NI4vrP|MG0?ieVee3w{I0w_?1lD}L#U&wUT5 zIk>;GdE>nO8AnMRhXAbd{IQ%WnX_fF2?-^k{ccnaywy-5bpF&+7xA4txno<39!)G)BP{7*9 ztzAYnzm@DpFG5z1lP0WRR6M5n9fQBl^D9xi$Yf;HC8^-6xclB)e;9An&2ZE%yu&Df z`(l8v3p@HN9OH9?!QZx55Df}9TR|52rOIX^GVASgzdWdldV4dz{&kre?`%t|smKE$ zW<`I!XFV8s8d;CXP0{Z&aZjmq4zc!Y4b&L+>PJ;c`cBsATHBg>ebm=<=Ri;m;t7LsI?IX?_0U zhW^xIcP`f%B=U<8C2Xmgj`wsvqbT3?nz4lqmGHXBOMmj)T{cwgEh_MlO>5iFBIO)o zAxPacSBE;&HJN6hH}nG)0a8BVwk>7MEFG(z1xwlvOnH#TqFngiqK-Rl^&c!#v}KGy zl3NgPSZ>7Y?>DP!T%LHarH}Ah&i<9|SB1K6=eASAckd(qkooVdeYl<#>C^i|jk2P%ONe0H z(nPX398F3K*p!QpN-ksrG~p?C-qGb_zLIq*WbGz?E-s~wj4722BQ>(}#P%g6qfH;* zdSTzv%+U^!shz4!IST%!_gaFXcSj{LTLt)V3|kr5gU>zm@|mgxycolgwsH~*?^jbl zKJ{8URwMAZyYBy&lnDAy?)0COsHDAv4*itg7&3W$Q#@i8` zQjG>S6>{x{u%(K`96)pQKLU^=%ne2rqH!>C*0!qR8GH8DXSf*;UH6LAu%8hQ)r zsLCMCCR9<{@u-A8FJLLmM7+O5_HqE{9w1U9W*bR-SwS+_sHW~GT^jn zqJS(*J{uQMd~@Rn#qsaWHY2bwmHR|$UPjhSz5swA5#LscQS)i?NJoB5>)0f97<$1* za2Cb!tO5mn9{GJVv?5^bd=$bGn>UjgDwccLjdqWU7Fmiz!_r_%lyUEQpLb}@{S!z^ z&{NRcPLxU;-1;R;{Crx+p^#7i4y~>c;{h-qA;PI#xgTon8Dq!JeiOjnLZy&2f42Ab zmwA(7AzeXC6ao_31U^J))Fst9ATb_w)`fm`9Bq>`_6GR?GXdP=RgTm8f9T`SDE1Q0 zjARk7-mdXth8nnWmDuUemd$DjZnjy%k>Gl;YUr0Q9YJ~@yU^v0XX`f>NKY^5fY$)- zBFqYD*PeKPbVma#AB{OLEu`1LQ-J1}{h==pk;70uxTo(MzoZ)G@4zi2ZF-d7wTJ6}wmeUeSi*?SuPXS5aTKw~s3IF4z6H6OCWU#X=dS0xGC>^O1+ zj~gdK)Iu{3J5eNEN6cpI7VD4L%!*B*F{U`=k{^Y?Iy}fO8O_O-tpM2tbrq+p0o_@F zBuZfOE#UQ3C)t`R8kf6(XBqP4Mu_^Gx%lZ{`Vint=pf~IvDauZ8Zhg%zWs9 z{547|2*l9W)hdeQM&6Jdq0-INPIZhBrb~A%ZJ`%g9mzTKjHAXeAp#y5^>qbTRjm>zT7d{d|gy9PHR!_&%aA)@Kbv z|I$Or4!KH@qPk^kyfXG2dasdk^F3Iye#^oM#1hW>(8Q52Uvl{h{JC8ch5n_z z6!|lVyOQ#I2(Rp>2TTxcOegcNFz^v2E7U$aYeR{)smPiiWVo4MO#pztrkQGu+3=aB z69`6LP@o%J`TCbGbtL)UeCC?XCu!EIzj%HXiUI4@t~ZW1)-OLsC7u#XET_-`(=mZ0{lsHs>UOHpLyv@rh&Vfg*f(9iF#!FFTPQ#v9q*T;26ZyX z1_4QH+cnqsq?C6@ExLhB-BY^{m(-`^$oM_{jnNaM19+q&mkti%dT_3;nd#UMa{P8?scyS^9*lR)cg z532%hY7!5DvzXEib4@pro>1njeXDt$Fn|S2*?mJxQIcfVG0urJswh9ol89p;w2;jl zAenRyd*q-!jIN%o&|nLL$iZh4RlpJ@CBOTZ?XkJ@o%ZlZV+$Iz<>r{&w*7IRj5`t8 z$Sxe=a@ljJLeKdgEt)*244u!0?^i@x3EuY$S}5;T+rtehO&o12Zh#oICCzU4Zzo(_ ziW$V97GJ`bm-#Wer)04&0(<+yzd|P2g_WSd(Py3rmaxORu#0(VFI+a4g<52O1yvMCdM5g-%IG`ur$I~_ypzJ=YJ z6!+r9z^kU+!E6~jUM-mMmzxN;=L6~nB8;N@8J=+wO!aO`FpWY0egcu#uaK zmAS(XcK*FZs~ZkBkUT)ab<)qjP>gTjWz6NCZ5V#F(5=!YZSj;|m}A`y#P($4iBh-HnTeoW0z zQXws!UKC1EUy^ZsvGSN*_F{Zh-XW}Z z_E~ZcaWXW1{f1CEB^x%A+_A1yQqOpu_|3AhyzKl)5bPP_+=+^CgwRvmKto;8Y2D&e z^VBRm>pN_{H*O7KV=6zFa#GBu9oRv^HCbBMHZK{WRfb;o@@# zOV1r+xh2$&ezdKN@O;9kL8gE`)3zVELh5}n#XGt|vs#u`_s6oVE9j5Z3==(4VS#6* z;_uEH8DdisGJi9s5aawI@IXIVycER& z0uez`mw8{P4g`kapAY8F>QUDv-ZJf;_5xfwuuCLVcvzAw4hR?Au<|ahuiy7*nH>^D zhUwlEo(MMr{7L;pOJ_~nx?wn$49>o?FhGm4DkbJj{cS*+<)mjV;E2|+%>c~Q zdRz~5A3r(<-;+bBV98f`g-5sbget)%HfIUQIvmx|ht~f58(CqoDWD9EkrZZT|Mn@7 zCHimreTdFK@J8!a%#ha!vhuLlWN}GyUN$0(rcBrpZ;H{Nm zmo~T}-!nUEhyR=a1Bj)uviX#VAq}lU2_l$_-~~xq4gWdTg57K2j%qeTW%$A>CZe?J z1vivi`(U3nEuLqTiw|F(80x>r((1e{&ARG9Y3wLf;=DDI&3-n1*ZFi}QI_?(1I1g>ZsL4rKe zyis;kO(2BqP8*N+W(n+Tik$Tpvzy(xNTexj`y`)PG^2wQLY=> zGwJutlU66l^Mfx0`|XVLL@IjSE_uZl!Uyd)-*HyQ-_Mgy(wM9;hVV zeI0kx}tKa5NB1FXhqwX zYF1&gpEOtHUlSBixAW`FOUKZjo>mwSJ24!wd~LsSQ^~nL%GFOd)_r;dSe@M|+8ZS+ z$}-oXsx722l`bQ^P54^(iR3*)IsSl6lZ?Mdf{3Pj*)fle8b=gm6qEv>yEgfyS}{D^v4lT}UF+*7(TBHxS1J{`Gj3 zH2!h1OOJL>si(F&T=&v84STNkOK@}8KeH>F_zS$!L%x^UYI)hbKVa=GX8(u%>f|#J z>Bckka(Oqm(SsErSk<_a$z5YzD?K*eP8U3#4C4>hn?ax2|oGWNOXT`T#A z@|fbMNq#pPZJyD%9@o7LI6*tW7BRJy?%n1%HeH?d<)o|zZv*gBZ}Rtm3(vqHy@1ft zYTU)>o{O{KEp&YizQ?R{{LKEkye9r`x~J|?qd z{k4!{H}k`RD-i%wJi^KSu8(#z1674|mOX=dCaZOopKvN0(L0ce?R-tUD592HZ_cgM z>a8W`eP~B(PNWrN{FM_zA08xv@JmZ+o#sRJ9c{7RUtRbD04pos=rSO4BZ3df;#dHZ ze^~h$nmWi(QYJbl6sNy~;Dk1r5E~j)+|WCsMp~Pc+kepM_8|`)+Ox^?)yJ>=cvzET z^6})c{=H;)J@3>LmVXyWBh$&FXj1($!*Cas=%Tn@Yr7JIV)8FhkR;50A35}Luhq9& zva6h^JXE>q^B=AeGAOU3!2DCg`l%PmRHHHeYw&Ld3@}kuZNfXy;FpJB;A85%;upGr zVS+!a-131awmXJoy(E4;Rvp!C={nt;==;$&f8RZqlZ6oawc9u-u6$zQNmt3a5jF<% zjNP&!v_>#wbw;^cG)j`n>#OLd>6}l6Zo${5iE1Ry%Jh8KxBNpLbGK z($Pd8LJAcFdZ`Sfoh&O)vuz&1|4LKm5%mx>%X`0(sB}8^DO-HTXg&`nL?(>gK5sE2 z!J~Ze1@uv-XHqV8nY_5W!^LHZx@!heu0u)RbI0KlI4e~~S}E(V2+(_W)DH{Nou zJ!i&x>Kj?;Hp1dYV$Lb&nGYZ?T2yUQsDXCi=pTKyVC3a&yfzsnPgvv+Eb5D)NIY$6 zT$vO*b8zmfy{OgB)9yDSA$j5%Iy%{JjNW9o6s+wMw7J;YD#|Mk8g0shNg?}a_m`s# z#e|GMLSeFHD>cWEC|-GK{`FJLJdijzP@u!0KMUH&W59xByZfLAod6^QEz>{1Q1c(K zPzGX3)#-Vu0f88Xn(DW+EVCsOpGpEW2zW>p346YV01ys2=)H=iD7Zimm}BX#X|1Kr z`jX_(6l4=Zbr|=AOYt%>7$K=dzux(UJV+8kiMP!=y(5j^NAd$Ql>I&7qB9YsD*fZG zhuTdjG2&G1vAIkgZeTIb(wlk|QhTWbGYMU>55}8;KuRIz8`OY2HJ^ynitl?D?iO)) zd9VIpxFSCi`(;0bs)cTb%7ta~}PnZ9a=v>7!Ju)yMf-Eqt`2ARl2VW>Hu)dMcXwH5~ zs}Xeuy8y={U)d`9+z|`Ikq}`vEia-8TLA8u* z3h5#^Yu98Hz+ihTieacEv|pG(Hp=*T zO^U#T3&4?K`1O(ss5FT)utHA;l}CODhR3eFfbI`SQb;?z5fxC)wRh?_(^BmG%6vpF zNRp_JjQAe#fo4!RX8yPb<9-h$deC(ktQm74CQt!2*h02xd%CbfQ%Ts78HC(Xz;TUp zF@|uNa8w`Tx}jQpT@k1J(CAbN@Z~!c48{@+ghX8Gb2m>oGp^pcI$S|uf(`0``U$u; zv_^5u%%JR?g)S9~9ys7v87j+;u(AL>r9K;F$)tU0K<6duNvLUGS_vS@fDA!~Fnb2q zNYG^;2ukH{Pj2gw4ns5LPx-Z(iV`0<$wmG&lo>4b$tnAu$HF{9ky%geAeI8_VmH?S z$%M;Jwc7z9yT6RJr`)M2_o3O1OYIZoA*~^=Mrj-jEUEAd9#!BWNb!2B^I3g}f#lh@ zu`hr`d4W-D!QaAlwCmSAG@Q~R6TMR-*^0E^M`86hln7Kf5bg}?gL(FM;_&m6&iNCe zn{S+s?R~vJPNQJzq6bX^|5S)3Or<>M^MGI4-o((KalpBuRJu#nS!Dgf13o9m=z)-I zEO&Ifbg(|Ny7UeHk)vCR`;wMl23hoTE$aEaV)s$lz2p0oT1mQ6Fe8h4Bq4iY>G3j z+(7VlnJ#~l6h-EZ+!BP4|Ib~|0p8MA^#2Y$CRiF*k39Kzesv`HezWf2y(QoV*g1*1 zUYz~U?xM12P*!lm?NC%6C?E0A7ktGK`V3Kjf3)Xf_;qLg`_`#$(Ka_{PX+fEJlY7k zA-_VC8NDAhJ`6m3$?!%9BlzL^3{Yzw)xSvua(h=2k5{(=v|mp^0@%+gMA9x45imZR zpMk~3($K-VkTv#YUF#3q%6UjG@|w{i#3o%^^eVH3;QHYDO-}GCNsAG!G1ShV1yT04 zq5PUEUC5pqz$GW7E)n6Wv?0t7dUFx9yQ;(!TgZ6DD5^~cQ|gxG^!jX^jnQFwFbpjX zEq4yEy-eMjeL+*a%}@G5&ZqHoP(|f23rx>-il#I!%IDi7p}zXVMRz6ykB5oja-RzK ztvPx4z#%zrcNc|!f7UMwj0Xw;wkL=2aF%6y5yBM$dh5CT%V6cyCY<7X7-yQ#m}7 zijc4Ue)?gPwYMzHkKg@(_)^$AJ-`XHtiB0o(N{;EOId_wDmqJfnn<@vPcQoY1ugF39rX)aJ8mNVljqFA6s+PDhvAUzqMXsgu;w4*Ay- zxrUD#v;RQ4j>E-TyQ+X=2R*rY8DMEDn|!0z&CHlyothu3#JIo=FhhPU7z%wKAw{Y3 zAUg&fT!c2WhSB`8akZIWGq@q|P>N~3vV+fz1fMQl4Vu2SUlGN&JJQ zr?XUv1eRheo=I}`jm^}%s9{wwD)CckS#9T#m~|D=)?lr!v2HDd%n}qcweOa1x0BX9 zLOHqywI{-Jfko!fT-ICizQ|)`=r_E0ut(~f^>>vGW?m@%bZ94j^-I)aG58V_zO?sZ z{G0hc*}Dgg@tO~zVUuZ3B^A{tJ3JG}kRXL_$)0h+W3jvUqNA&!grxfT241o^!`N#1 z5okxdm9ntz$+R)b!5cX0gkSvRBhdJ&H0pJ$G2S)aVECUD?sTL_q+6Wbr7BsCp_UPh zUay2CeviE3!C=xx;!<=;kamkDJN(X!DeCbfvU|k$dZB;KA%Af7;B0W&1SAN^)6^M9 z{O|WUw{T3xYmmdbFwZ48<8yB43#Hq;xc4N%Mjy?#eEF}3d73sgosUM-z>&Qh}eM9F3DbC6jyb(4ey>pYM z4+PnWYS#<$LZ3>tLO#aPnyb(F8b&2sx;}8sK%>^ft{0z$9Rodgjl3`@1ya6hU@w1} zql%T-k&Mz|SEjq<7?&P!gHZ*{$P3hm*GKuH-dwo|5r(r`^^+-DY018pQ7b)n`2;xb zl!(o|Lp@o&=mZ4#y{3rkGZ-8WN=_p`Ch)8_1n@A!d8{Zb9G^D$pp`V1uu}2@Q4DPl zM>U%tgN&LYaPXtm+q;|u19Bqn@q zenRq~wFW?@*S%#qLedfRSGODPvf|m`@u^t2O8yDc;`NdkP|5j|M!vf-M10`89jUk{ z<$ezIz5r_L^xSJcX?U~}I?vz1^wv24>Yi5ktKHcwuP49@7$r(dLQw5EXY@l(OrkjMrs@sGcA_=3**#ZXJ1ll{y|RSgXZT^*9_c) zoOh97cO$eM{8LwarL#DqYvv$L;A#xs3yKs{PZ@*Glv1pD2SzLwTTggS>z_hwxC*2_Rep=9sD>cnWY)60MG?p=A-7s6x@9 zUA@srK9~+*i;1XX6tBGboPWOJh>?My`JGigb^1V`G`fqYLYQzHO?J}X+bxZvgsy`Q z36J#Fj5%5h)=yVN{ssLfr1-Ub2(eH}hlEEjjh*xa45(`x_X5PEz5n2b@j&zKO|$Kk zWhNC4Bd(;Vw%=8OY6<1H6>2JoWouDyTwfpa*+-O8I;Y~O;))@SJrp7-VCMv5k9?3d zk^yYv>pSMp|A=v5&_S^Tm{=UP>gjOFwO}Z{LV8~SY;r6rbJ8QSuq{kts2T6Sq`Q55 z3vXoK2n+;=;!b8xL{+{*^Po_W?W7bzK#(%pY0sT-<)xy>1#kFcs!0qi$A+LS6qbum zZV!zqYb)3|l4u?}Pi))Q2}Sz&wd;`+(M?+68jjoK=F*53h?mIeup*NqJ2#U5Wt%I&ZgfAUCegcu31By_0oAhDKq_y z!L%}eaA$m!izb_89!ih%Z{XZbJ^;r%aV1~MQ^xrBbK9rjSQz*LPsX1f8bwD`8Nclv%Vi zLNP=#1v}7G7gmq+X-5zt^?vWisx!~?8^fft-jvfxhAnO*j)K`@+tGBBm6wmEFyriuK1Xtlm@zPu=0v2e1^6}^N_cl9|xjkbFyWb zmHoMX_QGMlfJyOhR`&4%4jG{Jxyo*1m2($=wnS-xUg9gxjG_+o&s@;@!u*{Hv`=K3 z{;=VuV2c(Y>dBbJ7k%l zD^a1liW@XT<6JR3Wk-9MHt=1!^|Lc#<@u?FR&2?onf@ujN^+RBD$uJA|F~#qp98+= z86Oj@SQPoVG+Y1XM{xZd{x9W`rdZYpTz#5lKxH>7n_k-6lymi%=2VmmfyYxK5+(<1 zhv#_W%oa%AsjqFtKR5P`JS(2572=@GkYk{~Qr=Bw7dwbJ+A9KQ9Fx~0C0a&UUIt-j zg&rh`zh(RM-F-Vvoj}qn#CK{PNvc=0tVS=qqe|HGYt`D*!0sor;dzvz?^L=QNC#u< z4qCA*gxf|W0iUXY_KuwNR)8T~lOhLPqIa5Nm4*yI!}w~9adyoK|q~2lx!&riS z&tw;@0~TCN7D8MNo8O^P6t_xR0c&@5sQ#ohEY5dBf`eN&anPWwXgjn7M|IZ8TU`eV zk+}M?6D8MUMa0h9@x+ape=2+m>Mv6D*QBK_Ql*=4U5PD6!ZgbibI1$jL@Ze#@NZsj zL(Ineqlvx8?Rv2IWHjOZeA)?dz=vuAfzb3uHCqb;$}=atZ@KTc=Bn#LD|<=%?NMPJ zMyI>`JP|yEQvEwdw(uKUq?$GWH3iYca|d5GObTv%R5rdicNV#c}C=46rsm zx0gC_gs9hPEQkU~1RIoLNmE`sY z5a%(}zE+T9JV$WvJ(h~Ms~xw6tfUA%#M0XU4Zbcl;Oc^bPVCG6yp&-cj*MAiLSI{I za3~om2^u@FKmN9i9(kD49liJ4-RgA$Ei&8BTn2yL>RIEL<&Gjv>OIQxG~CmbWPl!e zg`8@+92B3!EbB4gwLRO{m4}-R6JC~s4kz?=FhxE!1bg|rUB2Q6_VrG(q7G|?PNV{B zoY}w|*Xpi(l}juNljTgjEAQjx+7gH19Nti~fAuFbE=mJlF1oTpxMB{m%CuDn8+%`r zQDu2oQEnZs2r`M&y6tpS6%fi+V^)8l-A1%06XOWYi2dqPRNMN>7q}5*nNVn_5QWx# z-$JYZ_BUhpeNIjk2{m?&S{8z&X2k!h1W}QFj1BWrgyd9DyfK<;!+# zJq5|9Q7_2(jVXhR=c-Zi_pvc2Z?`z&44CKV4m;_L)Td4s5{*5`y|uVB;lW~m6^%-r zPJxQAR@x^?JTimTJmrFF@GnQmP&RF##WYq4AW6NP$jJ5Iju5PJN$Rn=bUN+g&%5k_ z;&f4ct?nWPthKF2?96@8HUs>2!7=P+I`ZQm^PhXSGxhBh2DjpzTCepl`?DEdJAt6L zzh~Z*%U>@10jjI&^F%JV;Xseh{|eJL^5!VW3ETH;9KbL#M}mxf6O^sS>tav0;| z*SWVMRIM$>T`jS)P^Iq5^&9m+KLPy}Ozv+E?IPANo3C7ACDn)+m~N9{)!&Su`U0aU zxAxoMBQTjYKoyAxHGm`3JR~`aDwx6!zu8Ps-gE+qg8t`P=owX%6_b-Oh1*XeozH$_ z_& zud$N|nnW=(7!d=7{UDr3XCI6cw09cWJNnl)GT~4mfX{)TEuc-YuvK6|E~w`FM}UTu z{|K(}-}CI6N!D*g>yGcYMKEpB&av_2K3+fqZ}}}RWDCuv%0JtP7dh4dUL`%(pQ1Gb zv&jMcAY=N&5ljlnrPACtWZjqPj&$wUJ)e9?VarS`r=)K8@8NX;^nay>wU?46h-zl5 z-|s<|Y;fqNdZnF2R-<{4ZWKYz^Y>d(b$JQ-bPlwisL{6UQJN?$KfG6?bJ?4x^j}ZGjxGkI~i!2=!6FzGMT~iPwQBTYr(ES zq6ox#{=wnbK1%odiECXYvx~Q6bIdK>p3~R*wNcjzgb#u*m9@CE@;SJ@odC1E|Nd$ku%> z1Dv@GS|ycf>6O}sHJwyMzphgf1SoOD+Wo$&Nl5i){%aye3D%BozQJ(?=r8S0io)xa z@2i9fo4DRz{`k6*UsSy*ZU)r#a>B=$qn88fwQcZ(dP~NmMX^4g7rge2K*eLyj_zb! zpbi!M*UQeH_9@&Yv(Q!$DedT~c~9IHXv@M`ER$GYIfm9^{MW3-Co50!**to6J4%7wxF^6Wjig zWzpyxU|(k8s}IyQiOsMUr1`&>56+zbh zg&zpBqzdm9OB71}{6bCHak`@C!>TPR-RFY4yeJh;x}wZYk>XoH6=&C8gE~BI7x;FEg1|ZZEn*u>25`EY9oBFBycm5#o9#sHi76!@cyzK+_bKs<*{@N)DF| z&ZsG5OS<#)XOdXFm5iqqVxemR^GxJkmlv}nU9-%42=f^UBV9CEPaQJ5*9hl~34Y6DXMGC_O%R%F}1re;ez zzx5_!XTiiX03S{B^=Xc(=6_qFevoC>2iln+M1)(G^ZPFQX+U{fsH#IS&5`Laf+<4I zJvF+aAc|`U1Zpu=>A&aXp6$^Qv_*f02=yk>k+e0$W5BQaQM*wLXtw6mjc8WRhJaa7mAK?jAr=golx6y>Z1bTzmv4xkhK z90~GE6y`|`poHuVaAo0D8v(Q&Xc?q%X4SqEN0=LrkO8%t%<(_|7@YHESw9fYY6ZP~ zy?BRm8Nw&vj1yq#(x5NFK@Di7{TR}qbUXw6N2?oK`F1ZSUN~S)9uY>Eq>8p6h%lq~ zJkvh+zT>e2&%82$FaBG|J$DC`7DmBAF3~5H6_n)cf$evI(~pN*E$3;$toFahAk7WX z653j~A>sY*wFV@0r8^sNop){G(XsIB)Ct(CfF9<gVI8TopQ=k zP<=e){cj05aR0E*Oa=mHmcoz->oym8 z@-O*u$hFr6lovub)2!pE#u;GoX_xP$T|U2$2|`b-)kW-Rv#~@-qRk zAC^bIPqfyhTVi#bIxVAJS;Y5rO$Z^JdM<=hfqrdFep!g~Z*CR5`e-=ni2SAn>*e`n zYpNx;%fw2g_R!F(8-U+M*HTA{zZL-JdkMG~xGd)ijuh(Vb_s%tM9Bpi(o|XQZ?_Cr zQ7$|ASG@cxjn~Rf8J$vs67DtEhg@<6l(!cNAC$%3ua{S9tV5|JM~#)LAzwHnV0t=` zD83${?K+j#mR?(E`wYwk8GfjdbS1EvHa@&c+7HB1{~N&n0k$IB%=cJ9TW)41;K?cK(h!tym#!>+;1zv>=w~6Gv^QK{& ztz`hJ$3Bf4>HsbUuy(x+@^;u}YnYyL=-C_Ff!h3z51>|sl}~p3Dh*h&NDdOM%|(m` zryTF)0_SomfJbvUQg-VFviCsEd9Ql0ZW}nW0YM2?uaVP=3=OCv_WuRIyKuJtN<%va zpc%XdaR{o_H&I%^xWd&c3?mQa%~jx)hzZa$1JK{MiZXm26RZVv<@wI&2<+CY-YWt3 z;?2*4c?{~k#}K;M`v%}8WZc-Nl@9=7n_>;{HG~07tn@E(&<$!p51Nn$UH+IM0KSV5 ze_QDf*(ABD7MCO23Qhc!;Fusp)J2^(@>4DP-AHu!p@RBh;K$Cjw&;WbJPDlNMdb{` z2UQo##K*(=sTBRL4Br^R@Iw3Ma^!@s%;Ec;E>>NHaF`>QU|+xp!}aa2MRtF>7Fqe} zGJ+bg3-DavsOW#NJ|9LXp|1yy0e(Y9^}kbS$G;MojS5IpA1w`Ea{KBc;9118I<2rT za2B#rjG7IDeiut|(*F1!jZ8duXvWh>jIj-d(Md}Dj{)z+iR~d)^);2eq<99RF<8G3 zS5jZiu~9!2s-bD-hXGGuTEHkI9U!y7TcZBcQ=zj!83e!*;Hy!-#smhRv=6QZwvPI! z@L4)7;B#IjEDdMU{yC30va=g0p-f;J_!lp4OZ`tOe5=55QQ!2%Kcl9=F35%hHV}q4 zX}{bH?B(=dw@@2%A+pxo`hM_z`zv5ey0rbf;TU05&?CUdqwhGMhT+4-!2L9y2Hg#6 z#B0spg`786h{xH`5-!bs5!gV6sJ}d+&Z8;B`)HD2_>%f&HX_rukj?9#1u8h+fb2Wh zDxkMY;^EJVCt2$8n_0q_fYG{JiH8mdkB zec)evt&E`%6vdD=D)2$zD@^bru|7hHe--cra_HMUrVTW%{0wBVP#I$a!-OIUQwQX? zfcr23P&Bl?zlP9S&+VND15hrs5;y^gHnXz?h6hFVgAT|?kX>jD#X&W$JCQ|sxtMf$ z1rCIWF zcK**3Y63n$hilOdnV}B%9|7McN7~9;I%q}D1};Rjf-d%qHROmTwBPxuSE!%iGGsQ$ z9D?~l&8!u`hZ%v{Hf1VHk)y>zL<$Xh{&dok&dc8iY#;SWQH1I47jd&!zNLKn@$aHG z&Q)j^VIQiqVkL5(z%AI!z#yPRxdp|q6bqxgR|D5L*9^TuY3XN(=fBrb?j{uFAkz4N zyO5Pif53bK{e}|z4oI}7d#%Xw{y*{l^oH)COmzitGP2{3Y6ZEWNDr|Nf@Q#4fj?nP zpx;mj-NAh*O9$rfkj*Kzd_#9o2JmS_i!M=sKdtB~qAH@c+ycBG*jK87%Tw|qlqP}TpZz?Tp~T1Z>=V~%(Zd4Zn*e}mYG<`d`-R4sodup=>%7Fpg5 zaFbWQAv;2ZE+0i=YSuf~b4R~|MARI(6nH)HF0`fz6bmK5-GQfh*NQCf_CAy~j&MEy z9|itM!F1BkEuNJL{0#UjWH%v;6DStCLaze6un*;_l)49zTy}(SXHeDun~2wc?4l}l zPdb1Df!`x9XmJ9?&m;dDnG)1~C<2=AQbY@I%eR!RpRNP8_1bf!kfqYmN$~88I6AZe zY*(55leXyw$B`zxbzjk&i^O8kEU4vpHcFWPw#g@hUInMDTSJ%~@>TTWff zIhZfh4A=sAMZsl-#M@Kd@zeBVTD)|qLS zoOFzT9ogJdw_Tq&=t;EzhXXet4P%2Rzod5NcPUo$)4UHS-WS69usJF5{}R|LYO{$X zeiu4JupP4Q*>5aM-6!{az=MgmcF!A{=kt**$=$LoNm6EfGq9i6ZW9k}h%)d#Vw1pT z3Gy3wXM1g(Csc~vwMN3Mr~!p@4>Z&o{oeO`-9k0=Skac8O7}pp!scu zgalikHMjq2;u#rRHQ55%IP-yj03Jsg(0l{kCvCSAqqfak`X~cfh%o>+xy1jvsEv&H zGr-jb4nqz^&_ASB-R>Oi3f7fu;)u$ngCT55S^&}6pZU*ITNZfv&GXmxV zXm00CzI2aCU562W;c!*grAX|+%`|q1ISZ4Mw#nZo zv}Lz&eE^fl0Q6ng0Dm2|g%N+j@Xdn7z&nAT6Z;6Xb4)8}9<0lN-N;7^a2deOUOD5Q z^gY*N#6J-Dj9?lGxVVTohCq{P1&xFD(@~rCss&6c1JHMTlz73F5&wXYUh=j4K0h!_mKkY`fk|yBYd~B(0!~m|4B15P){EwyFaEjF>y_yBkIXWn_a=x8tS}RqlPrdF(4l9suLW40K?h>aipcUn;gw`vEY&J?u%%aLW++Af{8 zBZw)X_VFmUiQCWak?Svw-ZSsTRE7aVO@t)rw;k{j;7laO;J-;HB-dmZB?8bv{vBd7 zppyZv?RJuv-zX=y53fXgNA9~C#u%y@_(V_vo`M{I{wd^G?1xAn@TwvB+QP6g2d&c_ z@P_DJI~nMz0esSV??Jy)FUhU{XA%EMr)^>w6R38Oq!v#BdjWq5oPw-WybhUa+)R6- zo1>9AxD2A-OhU{4BXX{wzI>QT#Z{YgG4hbXESxfOk0WHa5g>woCU2oTLNB$3nzx z*aO%b*e6Q+0ZWjnqxl)jOUggYP)GgLft_}b``x{E0dj7G5<|Q4LS!M(&w*zVpLaJ( zaO-?EvR2|I=i1op+iwgYU%14RMsOemvH&q5<^#J%|IY&s0k(|t*#luB0r9zh{lG@z z!$URzZv=ip`$JpwM4Nqx&$tU2t#PJ@*v~l?cfu=%Me^C zye3M^6Y|glB1*^to%smyKOLE(ROz!1VshwrH~d+zEkT&bJmAlWk0NyLpG6h|JxM&c zR{?iLzgLl|yg$=fo9u}&xluxmcrD2?BqrqY=r=p9%8)CRF_aQgo|q5pj4*-4Ip>`O zt^WsUcAyPwa>IE;WV&}(=Y2*=dohdwT;GDr7|LEMdG{tIA4$)oxYQx&v`=Dk|7GEfm(c;|rCVFJU-CfZcihjYv?>E>WFL>KhC_LJ8;t z443rxoqWGSc9K?1vaN3d*^miY{0Dl^r3j%7g9FfILjMh%jKo!dZP#TOB8>V#)cWs) zc+&^10k~&F%mQC0K8Y75Tk0DOxkCx+7=++d^&W(}?=qgNkhK}gpln{>gC7ZKF!X}w zAY-cE@&8KjO5nEu|AH)OUPtWY8=Dw@!{}fVBu;>=hWvl@*Z%h<{&L_Sku?YkQ}+$I zU=j$q8>pmz3h^$>MXUcwg#Rhv0HL;#6295^2E%}%O!q7FuKxS>eTW<`dOQ+fq$pzF z1cSlQA(Z%+(AyDUzJ0$5nNEB?u#Hn*D~WH~pTW>)sQ3SsbT0W-+PfyhBZww(Bob5L zmSdiMgJH;^()i0{Q~%Yr_qBxE5&M4cg#3-@4Td3v66Oxb2BK-|zY@LA999C~1Ktd5 z@8xTRZ!iowRQkV(=7USph`tKk2%JWI?x35W5xv1MBG9?_?-8s2Nkp$I{&~R5Xxu8? zq`ko~f>7zNd;j^We-*eD_!96MU}?hhHjOtVu;svpPN5R}^~Ce@8<0rAtB^xduLquR zo>78V6TvVfF!jsO33R+H2F^zg9{DfeTErQ++PT*y`mj&+GZ+j7!j1jkE8+J}4-zp9 a82. - -import asyncio -import json -import logging -import os -from typing import Any, Callable, Dict, List - -logger = logging.getLogger(__name__) - -# Registry of functions that can be scheduled -FUNCTION_MAP: Dict[str, Callable] = {} - -# Dictionary to manually track scheduled jobs by ID -scheduled_jobs: Dict[str, asyncio.TimerHandle] = {} - -# Path to the JSON file for job persistence TODO - pin this to the correct place -JOBS_FILE = os.path.join(os.getcwd(), "jobs.json") - - -def register_function(name: str, func: Callable): - """ - Register a function so it can be called by name later. - Example: - register_function("say_hello", say_hello) - """ - FUNCTION_MAP[name] = func - - -def load_jobs() -> List[Dict[str, Any]]: - """ - Load jobs from the JSON file, or return [] if none exist. - """ - if not os.path.exists(JOBS_FILE): - return [] - with open(JOBS_FILE, "r") as f: - return json.load(f) - - -def save_jobs(jobs: List[Dict[str, Any]]): - """ - Save jobs to the JSON file (overwrite). - """ - with open(JOBS_FILE, "w") as f: - json.dump(jobs, f, indent=4) - - -def cancel_job(job_id: str): - """ - Cancel a scheduled job by ID and remove it from the registry and persistence. - """ - handle = scheduled_jobs.pop(job_id, None) - if handle: - handle.cancel() - logger.info(f"Cancelled job '{job_id}'") - - jobs = [j for j in load_jobs() if j.get("id") != job_id] - save_jobs(jobs) - - -def run_once_job( - job_id: str, - func_name: str, - delay_seconds: float, - args=None, - kwargs=None, - persist=True, -): - """ - Schedule a job to run once after a delay (in seconds). - """ - args = args or [] - kwargs = kwargs or {} - - def job_wrapper(): - func = FUNCTION_MAP.get(func_name) - if func is None: - logger.error(f"Function '{func_name}' is not registered.") - return - func(*args, **kwargs) - cancel_job(job_id) - - loop = asyncio.get_event_loop() - handle = loop.call_later(delay_seconds, job_wrapper) - scheduled_jobs[job_id] = handle - - if persist: - jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append( - { - "id": job_id, - "type": "once", - "delay": delay_seconds, - "function": func_name, - "args": args, - "kwargs": kwargs, - } - ) - save_jobs(jobs) - logger.info( - f"Scheduled one-time job '{job_id}' to run in {delay_seconds} seconds." - ) - - -def recurring_job( - job_id: str, func_name: str, interval: float, args=None, kwargs=None, persist=True -): - """ - Schedule a recurring job. - """ - args = args or [] - kwargs = kwargs or {} - - def job_wrapper(): - func = FUNCTION_MAP.get(func_name) - if func is None: - logger.error(f"Function '{func_name}' is not registered.") - return - func(*args, **kwargs) - # Reschedule the job - handle = asyncio.get_event_loop().call_later(interval, job_wrapper) - scheduled_jobs[job_id] = handle - - cancel_job(job_id) - handle = asyncio.get_event_loop().call_later(interval, job_wrapper) - scheduled_jobs[job_id] = handle - - if persist: - jobs = [j for j in load_jobs() if j.get("id") != job_id] - jobs.append( - { - "id": job_id, - "type": "recurring", - "interval": interval, - "function": func_name, - "args": args, - "kwargs": kwargs, - } - ) - save_jobs(jobs) - logger.info(f"Scheduled recurring job '{job_id}' every {interval} seconds.") - - -def reload_jobs(): - """ - Reload jobs from JSON and reschedule them. - """ - jobs = load_jobs() - for job in jobs: - if job["type"] == "once": - run_once_job( - job["id"], - job["function"], - job["delay"], - job.get("args"), - job.get("kwargs"), - persist=False, - ) - elif job["type"] == "recurring": - recurring_job( - job["id"], - job["function"], - job["interval"], - job.get("args"), - job.get("kwargs"), - persist=False, - ) - - -async def start_scheduler(): - """ - Start the asynchronous scheduler loop. - - This function is a placeholder to keep the event loop alive. - Jobs are scheduled using asyncio.call_later and do not require polling. - """ - try: - await asyncio.Event().wait() - except asyncio.CancelledError: - logger.critical("Scheduler stopped.") - - """ - Start the asynchronous scheduler loop. - - This function is a placeholder for compatibility. Since we use asyncio.call_later, - jobs are scheduled directly on the event loop and no polling is required. - - Usage: - # In an async app (e.g., Textual) - asyncio.create_task(start_scheduler()) - - # Or in a standalone script - async def main(): - await start_scheduler() - - asyncio.run(main()) - """ - try: - while True: - await asyncio.sleep(3600) # Sleep indefinitely; jobs run via call_later - except asyncio.CancelledError: - logger.critical("Scheduler stopped.") diff --git a/README.md b/docs/README.md similarity index 100% rename from README.md rename to docs/README.md diff --git a/flows/localApproval.py b/flows/localApproval.py index 5cc69a2..da89c8c 100644 --- a/flows/localApproval.py +++ b/flows/localApproval.py @@ -1,311 +1,242 @@ -# Copyright (C) 2025 James Brotosky, Brandon Wickline -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. -# -# You should have received a copy of the GNU Affero General Public License -# along with this program. If not, see . +""" +This module handles the creation of local approval requests. +""" - -import datetime import logging import os -import re import time - -import dotenv -import numpy as np -import pandas as pd +from typing import List, Optional from models.agent import Agent -from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents +from services.agenthandler import moveAgentToRelatedPolicy, selectAgents from services.API import AirlockAPIWrapper -from utils.configmanager import get_protected_json, load_env, load_env_json -from utils.setup import get_base_directory +from utils.configmanager import get_protected_json from utils.utils import colorText, get_sanitized_input logger = logging.getLogger(__name__) -dotenv.load_dotenv() +class LocalApprovalRequestor: + """Handles creation of local approval requests in Loxide.""" -def getLocalApprovals(api: AirlockAPIWrapper): - base_dir = get_base_directory - result = api.otp_find_awaiting() - local_approval = pd.DataFrame(result["response"]["otpusage"]) - if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"): - previous_run = pd.read_parquet( - f"{base_dir}\\cache\\newest_local_approval.parquet" - ) - previous_run.to_parquet( - f"{base_dir}\\cache\\last_local_approval.parquet", index=False - ) - os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet") + def __init__(self, api: AirlockAPIWrapper, username: str = None): + """ + Initialize the local approval requestor. - # Only keep rows presumably created by the generate local approval function - local_approval = local_approval[ - local_approval["purpose"].str.startswith("🎫 Local Approval 🎫") - ] - - local_approval["batchid"] = local_approval["purpose"].apply( - lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1) - ) - - if not local_approval.empty: - local_approval.to_parquet( - f"{base_dir}\\cache\\newest_local_approval.parquet", index=False + Args: + api: AirlockAPIWrapper instance + username: Username creating the approvals (for tracking) + """ + self.api = api + self.policy_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + self.username = ( + username or os.getenv("USERNAME") or os.getenv("USER") or "unknown" ) - return local_approval + def create_local_approval( + self, agent_id: str, duration_minutes: int, batch_id: Optional[int] = None + ) -> bool: + """ + Create a single local approval request. + Args: + agent_id: Agent ID to create approval for + duration_minutes: Duration of approval in minutes + batch_id: Optional batch identifier (defaults to timestamp) -def scheduleAddingLAHashes(api: AirlockAPIWrapper): + Returns: + True if successful, False otherwise + """ + if batch_id is None: + batch_id = int(time.time()) - policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") - bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]") - pups = load_env_json("PUPS", "[]") - threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type=int) + purpose = ( + f"🎫 Local Approval 🎫 - {duration_minutes} mins - " + f"batch:{batch_id} Client:{agent_id} User:{self.username}" + ) - try: - register_function("add_hash", returnFromLocalApproval) - register_function("move_device", moveAgentToRelatedPolicy) - except Exception as e: - logger.warning(f"Failed to register functions: {e}") - return - - try: - approvals_df = getNewLocalApprovals(api) - if approvals_df.empty: - logger.debug("No new local approvals found. Nothing to schedule.") - return - batches = approvals_df.groupby("batchid") - except Exception as e: - logger.warning(f"Failed to retrieve or group local approvals: {e}") - return - - for batchid, batch_df in batches: try: - duration_minutes = int(batch_df["duration"].iloc[0]) - start_time = datetime.datetime.now() - run_time = start_time + datetime.timedelta(minutes=duration_minutes) - early_time = start_time + datetime.timedelta( - minutes=np.floor(duration_minutes * 0.95) + self.api.otp_generate(agent_id, duration_minutes, purpose) + logger.info( + f"Generated local approval for {agent_id}, batch {batch_id}, by {self.username}" ) - - early_timestamp = early_time.timestamp() - run_timestamp = run_time.timestamp() - - # Schedule add_hash job - try: - run_once_job( - f"add_hash_{batchid}", - "add_hash", - early_timestamp, - [ - api, - batch_df, - policy_relationship_map, - bad_publisher_list, - pups, - threat_tolerance_constant, - ], - None, - ) - logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}") - except Exception: - logger.debug("Failed to schedule add_hash for batch {batchid}: {e}") - - # Schedule move_device jobs - devices = batch_df["agentid"].drop_duplicates().tolist() - agents = [] - - for device in devices: - rows = api.agent_find_by_hostname(device).iterrows() - agents += [Agent(**row["data"]) for _, row in rows] - - for agent in agents: - try: - run_once_job( - f"move_device_{agent.hostame}_{batchid}", - "move_device", - run_timestamp, - [api, agent, policy_relationship_map], - "enforcement", - ) - - print( - f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}" - ) - except Exception as e: - print( - f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}" - ) - + return True except Exception as e: - logger.warning(f"Failed to process batch {batchid}: {e}") + logger.error(f"Failed to generate local approval for {agent_id}: {e}") + return False + def move_agent_to_audit(self, agent: Agent) -> bool: + """ + Move an agent to its corresponding audit policy. -def returnFromLocalApproval( - api, - device_df, - policy_relationship_map, - bad_publisher_list, - pups, - threat_tolerance_constant, -): - """ - # Get unique policy names from device list - policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist()) + Args: + agent: Agent object to move - # Create inverse map to go from Audit to Enforcement - inverse_map = {v: k for k, v in policy_relationship_map.items()} - - # Fetch all policies - all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()] - - # Define policy types - policy_types = [1, 2, 6, 7] - - #TODO finish logic for adding hashes - """ - working_dir = load_env("WORKING_DIR") - policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") - bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]") - pups = load_env_json("PUPS", "[]") - threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE") - print( - f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}" - ) - - -def moveToLocalApproval(api: AirlockAPIWrapper): - possible_durations = [15, 60, 360, 1440, 10080] - duration_selected = None - - print(colorText("Please select a duration:", "white")) - for i, option in enumerate(possible_durations, start=1): - print(f"{i}. {option}") - - try: - - choice = int(get_sanitized_input("Enter the number of your choice:")) - if 1 <= choice <= len(possible_durations): - duration_selected = possible_durations[choice - 1] - print(colorText(f"You selected: {duration_selected}", "yellow")) - logger.debug(f"You selected: {duration_selected}") - else: - print(colorText("❌ Invalid choice.", "red")) - logger.debug("Invalid Input") - return - except ValueError: - print(colorText("❌ Invalid input. Please enter a number.", "red")) - logger.debug("Invalid Input") - return - - agents = selectAgents(api) - batch = int(time.time()) - - if not agents: - print(colorText("❌ No agents found or error retrieving agents.", "red")) - logger.debug("No agents found or error retrieving agents") - return - - for agent in agents: + Returns: + True if successful, False otherwise + """ try: - addLocalApproval(api, batch, duration_selected, agent.agentid) - moveAgentToRelatedPolicy(api, agent, "audit") + moveAgentToRelatedPolicy(self.api, agent, "audit") + logger.info(f"Moved {agent.hostname} to audit policy") + return True except Exception as e: - print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red")) + logger.error(f"Failed to move {agent.hostname} to audit: {e}") + return False + def create_local_approval_batch( + self, + agents: List[Agent], + duration_minutes: int, + ) -> tuple[int, int, int]: + """ + Create local approvals for multiple agents and move them to audit. -def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid): + Args: + agents: List of Agent objects + duration_minutes: Duration of approval in minutes + db_path: Optional path to database for history tracking - purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}" - api.otp_generate(agentid, duration_selected, purpose) + Returns: + Tuple of (batch_id, success_count, failure_count) + """ + batch_id = int(time.time()) + success_count = 0 + failure_count = 0 + print(colorText(f"\nπŸ“¦ Processing batch {batch_id}...", "cyan")) + print(colorText(f"πŸ‘€ Requested by: {self.username}", "cyan")) + print( + colorText(f"πŸ“Š Moving {len(agents)} agent(s) to local approval\n", "cyan") + ) -def monitorAuditStatus(api: AirlockAPIWrapper): - current_agents = findAllAgents(api) - last_agents = [] - if not last_agents: - last_agents = current_agents - policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + for agent in agents: + try: + # Create local approval + approval_success = self.create_local_approval( + agent.agentid, duration_minutes, batch_id + ) - # Reverse map for audit β†’ enforcement - reverse_policy_map = {v: k for k, v in policy_relationship_map.items()} - known_transitions = set(policy_relationship_map.items()) | set( - reverse_policy_map.items() - ) + if not approval_success: + raise Exception("Failed to create local approval") - # Index last_agents by hostname for quick lookup - last_agent_map = {agent.hostname: agent for agent in last_agents} + # Move to audit policy + move_success = self.move_agent_to_audit(agent) - # Result buckets - newly_added = [] - same_policy = [] - moved_to_audit = [] - moved_to_enforcement = [] - unusual_move = [] + if not move_success: + raise Exception("Failed to move to audit policy") - for current in current_agents: - previous = last_agent_map.get(current.hostname) + print(colorText(f"βœ“ {agent.hostname}", "green")) + success_count += 1 - if not previous: - newly_added.append(current) - continue + except Exception as e: + print(colorText(f"βœ— {agent.hostname}: {e}", "red")) + logger.error(f"Error processing agent {agent.hostname}: {e}") + failure_count += 1 - if current.groupid == previous.groupid: - same_policy.append(current) - elif (previous.groupid, current.groupid) in known_transitions: - moved_to_audit.append(current) - elif (current.groupid, previous.groupid) in known_transitions: - moved_to_enforcement.append(current) - else: - unusual_move.append(current) + return batch_id, success_count, failure_count - # Return all five DataFrames - return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move + def interactive_local_approval(self): + """ + Interactive workflow to create local approvals for selected agents. + This prompts the user to select a duration and agents, then creates + the local approvals and moves agents to audit policies. + """ + # Duration options in minutes + duration_options = [ + (15, "15 minutes"), + (60, "1 hour"), + (360, "6 hours"), + (1440, "1 day"), + (10080, "1 week"), + ] -def getNewLocalApprovals(api: AirlockAPIWrapper): + # Display duration options + print(colorText("\n⏱️ Select Local Approval Duration:", "white")) + print(colorText("=" * 50, "white")) - working_dir = load_env("WORKING_DIR") - current_la = getLocalApprovals(api) + for i, (minutes, label) in enumerate(duration_options, start=1): + print(f" {i}. {label} ({minutes} minutes)") - # Load old approval list - old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet" - if os.path.exists(old_la_path): - old_la = pd.read_parquet(old_la_path) - else: - old_la = pd.DataFrame(columns=current_la.columns) + print(colorText("=" * 50, "white")) - # Create composite keys - current_la["key"] = ( - current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str) - ) - old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str) + # Get user selection + try: + choice = int(get_sanitized_input("\nEnter the number of your choice: ")) - # Find new entries - new_entries = current_la[~current_la["key"].isin(old_la["key"])] + if 1 <= choice <= len(duration_options): + duration_minutes, duration_label = duration_options[choice - 1] + print(colorText(f"βœ“ Selected: {duration_label}", "green")) + logger.info(f"User selected duration: {duration_minutes} minutes") + else: + print(colorText("❌ Invalid choice.", "red")) + logger.warning("Invalid duration choice") + return - # Convert 'granted' to datetime and filter by last 10 minutes - new_entries["granted"] = pd.to_datetime( - new_entries["granted"], utc=True, errors="coerce" - ) - ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta( - minutes=10 - ) - recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago] + except ValueError: + print(colorText("❌ Invalid input. Please enter a number.", "red")) + logger.warning("Invalid input for duration selection") + return - # Save current approvals for next run - current_la.drop(columns=["key"], inplace=True) - current_la.to_parquet(old_la_path, index=False) + # Select agents + print(colorText("\n🎯 Select Agents for Local Approval:", "white")) + agents = selectAgents(self.api) - return recent_entries + if not agents: + print(colorText("❌ No agents found or error retrieving agents.", "red")) + logger.warning("No agents selected or error retrieving agents") + return + + # Confirm with user + print(colorText("\nπŸ“‹ Summary:", "cyan")) + print(colorText(f" Duration: {duration_label}", "white")) + print(colorText(f" Agents: {len(agents)}", "white")) + + confirm = get_sanitized_input("\nProceed? (y/n): ").lower() + + if confirm != "y": + print(colorText("❌ Operation cancelled.", "yellow")) + return + + # Process the batch + batch_id, success_count, failure_count = self.create_local_approval_batch( + agents, duration_minutes + ) + + # Display summary + self._display_summary(batch_id, duration_label, success_count, failure_count) + + def _display_summary( + self, batch_id: int, duration_label: str, success_count: int, failure_count: int + ): + """ + Display operation summary. + + Args: + batch_id: Batch identifier + duration_label: Human-readable duration + success_count: Number of successful operations + failure_count: Number of failed operations + """ + print(colorText(f"\n{'=' * 60}", "white")) + print(colorText("πŸ“Š Local Approval Summary", "cyan")) + print(colorText("=" * 60, "white")) + + print(colorText(f"βœ“ Successfully processed: {success_count}", "green")) + + if failure_count > 0: + print(colorText(f"βœ— Failed: {failure_count}", "red")) + + print(colorText(f"\nπŸ“¦ Batch ID: {batch_id}", "cyan")) + print(colorText(f"⏱️ Duration: {duration_label}", "cyan")) + + print(colorText("=" * 60, "white")) + print(colorText("\nπŸ’‘ Next Steps:", "yellow")) + print(colorText(" β€’ Agents have been moved to audit policies", "white")) + print(colorText(" β€’ Local approvals are active", "white")) + print( + colorText( + f" β€’ Agents will return to enforcement after {duration_label}", "white" + ) + ) + print(colorText("=" * 60 + "\n", "white")) diff --git a/screens/moveagentworkflowscreen.py b/screens/moveagentworkflowscreen.py new file mode 100644 index 0000000..7199f55 --- /dev/null +++ b/screens/moveagentworkflowscreen.py @@ -0,0 +1,61 @@ +from typing import List + +from textual.app import ComposeResult +from textual.screen import Screen + +from models.agent import Agent +from widgets.agentmoveoperations import AgentMoveOperations +from widgets.multiagentselector import MultiAgentSelector +from widgets.resultsdisplay import ResultsDisplay + + +class MoveAgentWorkflowScreen(Screen): + """Screen that handles the agent movement workflow.""" + + def __init__(self, all_agents: List[Agent]): + super().__init__() + self.all_agents = all_agents + self.selected_agents = None + + def compose(self) -> ComposeResult: + """Start with the multi-agent selector.""" + yield MultiAgentSelector(self.all_agents) + + def on_multi_agent_selector_agents_selected( + self, message: MultiAgentSelector.AgentsSelected + ) -> None: + """Handle selected agents - switch to operations screen.""" + self.selected_agents = message.selected_agents + + # Remove the MultiAgentSelector + selector = self.query_one(MultiAgentSelector) + selector.remove() + + # Mount the AgentMoveOperations with the selected Agent objects + self.mount(AgentMoveOperations(self.selected_agents)) + + def on_agent_move_operations_operation_complete( + self, message: AgentMoveOperations.OperationComplete + ) -> None: + """Handle completion of move operation - transition to results screen.""" + # Format successful results + success_lines = [] + for agent, result in message.successful: + success_lines.append(f"βœ“ {agent.hostname}") + + # Format unsuccessful results + failure_lines = [] + for agent, error in message.unsuccessful: + failure_lines.append(f"βœ— {agent.hostname}: {error}") + + successful_text = "\n".join(success_lines) if success_lines else "(none)" + unsuccessful_text = "\n".join(failure_lines) if failure_lines else "(none)" + + # Remove the operations widget + ops_widget = self.query_one(AgentMoveOperations) + ops_widget.remove() + + # Mount the results display + self.mount( + ResultsDisplay(message.operation, successful_text, unsuccessful_text) + ) diff --git a/screens/policyselectorscreen.py b/screens/policyselectorscreen.py new file mode 100644 index 0000000..5937860 --- /dev/null +++ b/screens/policyselectorscreen.py @@ -0,0 +1,91 @@ +# Copyright (C) 2025 James Brotosky, Brandon Wickline +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +""" +Policy Selector Screen Module + +Provides a Textual Screen wrapper for the PolicySelector widget that manages +the policy selection workflow. +""" + +import logging + +from textual.app import ComposeResult +from textual.screen import Screen + +from widgets.policyselector import PolicySelector + +logger = logging.getLogger(__name__) + + +class PolicySelectorScreen(Screen): + """ + A Textual Screen for policy selection in agent move operations. + + This screen wraps the PolicySelector widget and manages the workflow + of selecting a target policy for bulk agent movements. + + Attributes: + policies: List of available policies (Policy objects or DataFrame). + agent_move_operations: Reference to the parent AgentMoveOperations widget. + """ + + CSS = """ + Screen { + layout: vertical; + background: $surface; + } + """ + + def __init__( + self, + policies, + agent_move_operations=None, + ): + """ + Initialize the PolicySelectorScreen. + + Args: + policies: List of available policies to display. + agent_move_operations: Reference to parent AgentMoveOperations widget. + Used to call back when policy selection is confirmed. + """ + super().__init__() + self.policies = policies + self.agent_move_operations = agent_move_operations + + def compose(self) -> ComposeResult: + """Create the PolicySelector widget.""" + yield PolicySelector(self.policies) + + def on_policy_selector_policy_selected( + self, message: PolicySelector.PolicySelected + ) -> None: + """ + Handle policy selection from the PolicySelector widget. + + When a policy is selected, this handler: + 1. Closes the selector screen + 2. Calls the parent AgentMoveOperations to execute the move + + Args: + message (PolicySelector.PolicySelected): Contains the selected policy. + """ + # Pop this screen to return to AgentMoveOperations + self.app.pop_screen() + + # Call parent widget's method to execute the move + if self.agent_move_operations: + self.agent_move_operations._execute_move_to_policy(message.policy) diff --git a/utils/selector.py b/utils/selector.py index 0a0f498..2cac574 100644 --- a/utils/selector.py +++ b/utils/selector.py @@ -33,7 +33,7 @@ class Selector: def _display_choices( items: List[Any], label_func: Callable[[Any], str], - num_columns: int = 4, + num_columns: int = 3, header: str = "Available Choices:", ) -> None: # Force single column if items are DataFrame rows @@ -54,7 +54,7 @@ class Selector: @staticmethod def _display_selected_items( - selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 4 + selected: List[Any], label_func: Callable[[Any], str], num_columns: int = 3 ) -> None: print(colorText("\nCurrent selections:", "cyan")) if not selected: @@ -93,7 +93,7 @@ class Selector: allow_multiple: bool = False, prompt_each: bool = False, header: str = "Available Choices:", - num_columns: int = 4, + num_columns: int = 3, ) -> Union[Optional[Any], List[Any]]: if not items: logger.warning("No items available for selection.") diff --git a/utils/tui.py b/utils/tui.py index 6eb070d..9fb1c54 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -23,6 +23,7 @@ from flows.prepPolicy import menu_policy_enforce from flows.quietAgent import findQuietAgents from models.agent import Agent from models.policy import Policy +from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen from screens.otpworkflowscreen import OTPWorkflowScreen from services.agenthandler import findAgents, moveAgents, toggleEnforcement from services.API import AirlockAPIWrapper @@ -30,9 +31,12 @@ from services.policyhandler import confirmUpdateAfromE from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config from utils.utils import open_directory +from widgets.agentmoveoperations import AgentMoveOperations from widgets.multiagentselector import MultiAgentSelector from widgets.OTP_generate import OTPGenerator from widgets.policytreewidget import PolicyTreeWidget +from widgets.resultsdisplay import ResultsDisplay +from widgets.retro_terminal_theme import get_retro_terminal_theme from widgets.themeselector import ThemeSelector dotenv.load_dotenv() @@ -106,6 +110,7 @@ class MainMenuScreen(Screen): ("πŸ”‡ - Find Quiet Hosts", "find_quiet_button"), ], "move": [ + ("πŸ”„ - Move Agent Workflow", "move_agent_workflow_button"), ("βœ… - Move to local approval", "move_local_button"), ("πŸ”„ - Move to Audit/Enforcement", "move_audit_button"), ("πŸ”€ - Move - Other", "move_other_button"), @@ -261,6 +266,47 @@ class MainMenuScreen(Screen): self.app.exit() + def on_agent_move_operations_operation_complete( + self, message: AgentMoveOperations.OperationComplete + ) -> None: + """Handle completion of agent move operation - show results.""" + logger.info( + "Agent move operation completed: %s, %d successful, %d unsuccessful", + message.operation, + len(message.successful), + len(message.unsuccessful), + ) + + # Format results for display + successful_text = "\n".join( + [f"{agent.hostname}" for agent, _ in message.successful] + ) + unsuccessful_text = "\n".join( + [f"{agent.hostname}: {error}" for agent, error in message.unsuccessful] + ) + + # Remove the operations widget + try: + ops_widget = self.query_one(AgentMoveOperations) + ops_widget.remove() + except Exception: + pass + + # Show results + self.query_one("#content", Vertical).mount( + ResultsDisplay(message.operation, successful_text, unsuccessful_text) + ) + + def on_results_display_go_back(self, message: ResultsDisplay.GoBack) -> None: + """Handle back button from results display.""" + try: + results_widget = self.query_one(ResultsDisplay) + results_widget.remove() + except Exception: + pass + # Return to main menu + self.app.pop_screen() + def on_directory_tree_file_selected( self, event: DirectoryTree.FileSelected ) -> None: @@ -282,6 +328,11 @@ class MainMenuScreen(Screen): _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) case "find_quiet_button": _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) + case "move_agent_workflow_button": + # Push Move Agent workflow screen + self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices)) + event.stop() + return # Don't exit the app case "move_local_button": _PENDING_JOB = ( "legacy", @@ -349,12 +400,21 @@ class Loxide(App): self.devices = [ Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows() ] + + # Enrich agents with policy information + if self.policies and self.devices: + for agent in self.devices: + agent.enrich_with_policies(self.policies) + logger.debug( + f"Enriched {len(self.devices)} agents with policy information" + ) except Exception as exc: logger.error("Failed to load policies/devices: %s", exc) self.policies = None self.devices = None def on_mount(self, api: AirlockAPIWrapper) -> None: + self.register_theme(get_retro_terminal_theme()) self.theme = self._textual_theme self.push_screen(MainMenuScreen(api)) diff --git a/utils/utils.py b/utils/utils.py index 26b2023..b6ed034 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -151,44 +151,6 @@ def irtang(): ) -def displayIntro(): - - print( - colorText( - r""" - _____ .__ .__ __ ___________ .__ - / _ \ |__|______| | ____ ____ | | __ \__ ___/___ ____ | | ______ - / /_\ \| \_ __ \ | / _ \_/ ___\| |/ / | | / _ \ / _ \| | / ___/ -/ | \ || | \/ |_( <_> ) \___| < | |( <_> | <_> ) |__\___ \ -\____|__ /__||__| |____/\____/ \___ >__|_ \ |____| \____/ \____/|____/____ > - \/ \/ \/ \/ -""", - "cyan", - ) - ) - - -def welcome(): - print( - colorText( - "=================================================================================", - "cyan", - ) - ) - print( - colorText( - "======================== Welcome to the Airlock API Tool ========================", - "cyan", - ) - ) - print( - colorText( - "=================================================================================", - "cyan", - ) - ) - - def section_header(title): print( colorText( diff --git a/widgets/OTP_generate.py b/widgets/OTP_generate.py index 6aff859..a3476c6 100644 --- a/widgets/OTP_generate.py +++ b/widgets/OTP_generate.py @@ -6,7 +6,16 @@ from textual.css.query import NoMatches from textual.message import Message from textual.reactive import reactive from textual.widget import Widget -from textual.widgets import Button, Input, RadioButton, RadioSet, Static, TextArea +from textual.widgets import ( + Button, + Footer, + Header, + Input, + RadioButton, + RadioSet, + Static, + TextArea, +) from models.agent import Agent @@ -70,6 +79,7 @@ class OTPGenerator(Widget): pass def compose(self): + yield Header(show_clock=True, icon="βš™") title_text = Static( f"🎫 Generate One Time Passes for {len(self.devices)} device(s)", id="otpgen_title", @@ -165,6 +175,7 @@ class OTPGenerator(Widget): copy_button.styles.margin = (1, 0, 0, 0) copy_button.styles.display = "none" yield copy_button + yield Footer() def on_mount(self) -> None: """Set initial button state.""" diff --git a/widgets/agentmoveoperations.py b/widgets/agentmoveoperations.py new file mode 100644 index 0000000..1e55fdf --- /dev/null +++ b/widgets/agentmoveoperations.py @@ -0,0 +1,704 @@ +""" +Agent Move Operations Widget Module + +This module provides a Textual-based UI widget for performing bulk operations on +agent devices in the Airlock system. It allows users to: +- View selected agents and their current policy assignments +- Move agents to local approval mode with OTP enforcement +- Toggle agents between audit and enforcement policy modes +- Select and move agents to alternate policies (future implementation) + +The widget tracks operation state, manages button availability, and displays +results with success/failure summaries that can be copied to clipboard. + +Dependencies: + - textual: TUI framework for building the widget and UI components + - models.agent: Agent model class + - services.agenthandler: Core agent operation functions + - flows.localApproval: Local approval workflow handling +""" + +import logging +from typing import List + +from textual.containers import Horizontal, Vertical +from textual.css.query import NoMatches +from textual.message import Message +from textual.reactive import reactive +from textual.widget import Widget +from textual.widgets import Button, DataTable, Header, Static, TextArea + +from models.agent import Agent +from screens.policyselectorscreen import PolicySelectorScreen + +logger = logging.getLogger(__name__) + + +class AgentMoveOperations(Widget): + """ + A Textual widget for managing bulk agent operations and policy migrations. + + This widget provides a comprehensive UI for performing operations on multiple + selected agents. It displays the list of target agents and provides buttons to + trigger various bulk operations like toggling policy modes or enabling local approval. + + The widget manages its own state through reactive properties and provides real-time + feedback on operation progress and results. Operations are executed sequentially + per agent with error handling that tracks both successful and failed operations. + + Attributes: + operation_in_progress (reactive[bool]): Tracks whether an operation is currently + executing. Used to disable buttons during execution. + selected_operation (reactive[str]): Tracks which operation type is currently + selected or in progress (e.g., "local_approval", "toggle_enforcement"). + + Example: + ```python + agents = [agent1, agent2, agent3] + widget = AgentMoveOperations(agents) + ``` + """ + + # Reactive property to track if an operation is in progress + operation_in_progress = reactive(False) + # Tracks the currently selected operation type + selected_operation = reactive("") + + class OperationComplete(Message): + """ + Message posted when a bulk operation completes. + + This message is broadcast to parent widgets/screens to notify them of + operation completion along with detailed results. It contains the list + of agents that were processed and the outcome for each. + + Attributes: + operation (str): Name of the operation that completed (e.g., "Local Approval Mode"). + agents (List[Agent]): List of all agents that were targeted by the operation. + successful (List[tuple]): List of (Agent, result_data) tuples for successfully + processed agents. Result data varies by operation type. + unsuccessful (List[tuple]): List of (Agent, error_message) tuples for agents + where the operation failed. Error message is a string explaining the failure. + """ + + def __init__( + self, + operation: str, + agents: List[Agent], + successful: List[tuple], + unsuccessful: List[tuple], + ): + super().__init__() + self.operation = operation + self.agents = agents + self.successful = successful # List of (agent, result) tuples + self.unsuccessful = unsuccessful # List of (agent, error) tuples + + def __init__(self, agents: List[Agent]): + """ + Initialize the AgentMoveOperations widget. + + Args: + agents (List[Agent]): List of Agent objects to perform operations on. + These agents will be displayed in the widget's agent table. + """ + super().__init__() + self.agents = agents + + def watch_operation_in_progress(self, old_value: bool, new_value: bool) -> None: + """ + React to changes in the operation_in_progress reactive property. + + This is called automatically by Textual when operation_in_progress changes. + It updates the button states to reflect whether an operation is running. + + Args: + old_value (bool): Previous value of operation_in_progress. + new_value (bool): New value of operation_in_progress. + """ + self._update_button_states() + + def _update_button_states(self) -> None: + """ + Update the enabled/disabled state of operation buttons based on current status. + + This method implements the following logic: + - If an operation is in progress: disable all buttons + - If an operation is selected: disable only that operation's button + - If no operation is selected: enable all buttons + + The state transitions prevent users from starting multiple operations + simultaneously and provide visual feedback on which operation is active. + + Handles NoMatches exceptions gracefully in case buttons are not yet rendered. + """ + try: + local_approval_btn = self.query_one("#local_approval_btn", Button) + toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button) + other_policy_btn = self.query_one("#other_policy_btn", Button) + + # If operation in progress, disable all + if self.operation_in_progress: + local_approval_btn.disabled = True + toggle_enforcement_btn.disabled = True + other_policy_btn.disabled = True + else: + # If an operation was selected, keep it disabled, enable others + if self.selected_operation: + local_approval_btn.disabled = ( + self.selected_operation == "local_approval" + ) + toggle_enforcement_btn.disabled = ( + self.selected_operation == "toggle_enforcement" + ) + other_policy_btn.disabled = ( + self.selected_operation == "other_policy" + ) + else: + # Enable all buttons + local_approval_btn.disabled = False + toggle_enforcement_btn.disabled = False + other_policy_btn.disabled = False + + except NoMatches: + pass + + def _display_results( + self, operation_name: str, successful: list, unsuccessful: list + ) -> None: + """ + Display operation results in the results text area. + + Formats the results into a human-readable summary including: + - Operation name and separator + - List of successful operations with agent hostnames + - List of failed operations with agent hostnames and error messages + - Summary statistics (total successful/failed count) + + The results are displayed in the results_text TextArea widget and the + results container is made visible after being initially hidden. + + Args: + operation_name (str): Human-readable name of the operation (e.g., "Local Approval Mode"). + successful (list): List of (Agent, result_data) tuples for successful operations. + unsuccessful (list): List of (Agent, error_message) tuples for failed operations. + """ + try: + # Build results text + results_lines = [ + f"Operation: {operation_name}", + f"{'=' * 50}", + "", + f"Òœ… Successful ({len(successful)}):", + ] + + if successful: + for agent, result in successful: + results_lines.append(f" Ò€’ {agent.hostname}") + else: + results_lines.append(" (none)") + + results_lines.append("") + results_lines.append(f"ҝŒ Failed ({len(unsuccessful)}):") + + if unsuccessful: + for agent, error in unsuccessful: + results_lines.append(f" Ò€’ {agent.hostname}: {error}") + else: + results_lines.append(" (none)") + + results_lines.append("") + results_lines.append(f"{'=' * 50}") + results_lines.append( + f"Total: {len(successful)} successful, {len(unsuccessful)} failed" + ) + + results_text_widget = self.query_one("#results_text", TextArea) + results_text_widget.text = "\n".join(results_lines) + + # Show results container + results_container = self.query_one("#results_container", Vertical) + results_container.styles.display = "block" + + except Exception as e: + logger.error(f"Error displaying results: {e}") + + def compose(self): + """ + Build the UI layout for the AgentMoveOperations widget. + + This method is called by Textual to create the widget's UI structure. + It builds a two-column layout with: + - Left side: Agent table showing selected agents and their current policies + - Right side: Operation buttons and results display area + - Bottom: Navigation buttons (Back, Reset) + + The layout is responsive with: + - Agent table: 2/3 width + - Operations panel: 1/3 width + - Results area: Initially hidden, shown after operation completion + """ + yield Header(show_clock=True, icon="βš™") + title_text = Static( + f"↔️ Move Agent Operations - {len(self.agents)} device(s) selected", + id="move_ops_title", + ) + title_text.styles.margin = (0, 0, 1, 0) + yield title_text + + with Horizontal() as main_layout: + main_layout.styles.height = "auto" + + # Left side - Agent list + with Vertical() as left_side: + left_side.styles.width = "2fr" + left_side.styles.height = "auto" + + agents_label = Static("Selected Agents:") + agents_label.styles.margin = (0, 0, 0, 0) + yield agents_label + + # Create a DataTable to show agents with their current policies + agent_table = DataTable(id="agent_table") + agent_table.styles.height = "1fr" + agent_table.styles.margin = (1, 0, 1, 0) + yield agent_table + + # Right side - Operation buttons + with Vertical() as right_side: + right_side.styles.width = "1fr" + right_side.styles.height = "auto" + + operations_label = Static("Operations:") + operations_label.styles.margin = (0, 0, 1, 0) + yield operations_label + + # Operation buttons + local_approval_btn = Button( + "βœ… Local Approval Mode", id="local_approval_btn" + ) + local_approval_btn.styles.width = "100%" + local_approval_btn.styles.margin = (0, 0, 1, 0) + yield local_approval_btn + + toggle_enforcement_btn = Button( + "πŸ”„ Toggle Audit/Enforcement", id="toggle_enforcement_btn" + ) + toggle_enforcement_btn.styles.width = "100%" + toggle_enforcement_btn.styles.margin = (0, 0, 1, 0) + yield toggle_enforcement_btn + + other_policy_btn = Button( + "πŸ”€ Move to Other Policy", id="other_policy_btn" + ) + other_policy_btn.styles.width = "100%" + other_policy_btn.styles.margin = (0, 0, 1, 0) + yield other_policy_btn + + # Status label + status_label = Static("", id="status_label") + status_label.styles.margin = (2, 0, 0, 0) + yield status_label + + # Results display area (initially hidden) + with Vertical(id="results_container") as results_container: + results_container.styles.height = "auto" + results_container.styles.margin = (1, 0, 0, 0) + results_container.styles.display = "none" + + results_label = Static("Γ°ΕΈβ€œΕ  Results:", id="results_label") + results_label.styles.margin = (0, 0, 0, 0) + yield results_label + + results_text = TextArea(id="results_text", read_only=True) + results_text.styles.height = 15 + results_text.styles.margin = (0, 0, 1, 0) + yield results_text + + copy_results_btn = Button( + "Γ°ΕΈβ€œβ€Ή Copy Results to Clipboard", id="copy_results_btn" + ) + copy_results_btn.styles.width = "100%" + yield copy_results_btn + + # Bottom buttons + with Horizontal() as button_row: + button_row.styles.height = "auto" + button_row.styles.margin = (1, 0, 0, 0) + + back_button = Button("Ò† Back", id="back_button") + back_button.styles.width = "1fr" + yield back_button + + reset_button = Button("Γ°ΕΈβ€β€ž Reset Selection", id="reset_button") + reset_button.styles.width = "1fr" + yield reset_button + + def on_mount(self) -> None: + """ + Initialize widget after it has been mounted on the screen. + + This Textual lifecycle method is called after the widget is added to the DOM. + It performs initialization tasks: + - Populates the agent table with columns for Hostname, Policy, and Status + - Adds rows to the table for each agent in self.agents + - Initializes button states based on current widget state + + The agent table displays agent.hostname, agent.groupname (or "Unknown"), + and agent.status_text (or "Unknown") for each agent. + """ + table = self.query_one("#agent_table", DataTable) + table.add_columns("Hostname", "Current Policy", "Status") + + for agent in self.agents: + table.add_row( + agent.hostname, + agent.groupname or "Unknown", + agent.status_text or "Unknown", + ) + + self._update_button_states() + + def on_button_pressed(self, event: Button.Pressed): + """ + Handle button press events from the widget. + + This Textual event handler routes button presses to appropriate actions: + - back_button: Pop this screen (return to parent) + - reset_button: Clear operation state and hide results + - copy_results_btn: Copy results text to clipboard (requires pyperclip) + - local_approval_btn: Start local approval operation + - toggle_enforcement_btn: Start toggle audit/enforcement operation + - other_policy_btn: Start move to other policy operation + + After handling, event.stop() is called to prevent event propagation. + + Args: + event (Button.Pressed): The button press event containing the button reference. + """ + + btn_id = event.button.id + + if btn_id == "back_button": + self.app.pop_screen() + event.stop() + + elif btn_id == "reset_button": + # Reset operation selection + self.selected_operation = "" + self.operation_in_progress = False + status_label = self.query_one("#status_label", Static) + status_label.update("") + # Hide results + try: + results_container = self.query_one("#results_container", Vertical) + results_container.styles.display = "none" + except NoMatches: + pass + event.stop() + + elif btn_id == "copy_results_btn": + try: + results_text = self.query_one("#results_text", TextArea) + import pyperclip + + pyperclip.copy(results_text.text) + self.app.notify( + "Òœ… Results copied to clipboard!", + severity="information", + timeout=2, + ) + except ImportError: + self.app.notify( + "Òő ï¸ pyperclip not installed. Run: pip install pyperclip", + severity="warning", + ) + except Exception as e: + self.app.notify(f"ҝŒ Failed to copy: {str(e)}", severity="error") + event.stop() + + elif btn_id == "local_approval_btn": + self._start_local_approval_operation() + event.stop() + + elif btn_id == "toggle_enforcement_btn": + self._start_toggle_enforcement_operation() + event.stop() + + elif btn_id == "other_policy_btn": + self._start_other_policy_operation() + event.stop() + + def _start_local_approval_operation(self) -> None: + """ + Execute the local approval mode operation on all selected agents. + + This operation performs the following steps for each agent: + 1. Generate a unique batch ID (current Unix timestamp) + 2. Create a local approval OTP with default duration of 360 minutes (6 hours) + 3. Move the agent to its related audit policy mode + + The operation: + - Sets operation state flags (selected_operation, operation_in_progress) + - Updates the status label with progress indicator + - Iterates through all agents, tracking successful and unsuccessful operations + - Displays formatted results via _display_results() + - Posts an OperationComplete message for parent widget handling + + Agents that fail are logged and added to the unsuccessful list with error details. + The operation completes and returns to a non-busy state regardless of individual + agent success/failure. + + Note: The OTP duration (360 minutes) is currently hardcoded and could be + made configurable in future versions. + """ + self.selected_operation = "local_approval" + self.operation_in_progress = True + + status_label = self.query_one("#status_label", Static) + status_label.update("ҏ³ Moving agents to local approval...") + + # Get API from app + api = self.app.api + + successful = [] + unsuccessful = [] + + try: + import time + + from services.agenthandler import moveAgentToRelatedPolicy + + # Generate batch ID + batch = int(time.time()) + duration = 360 # Default 6 hours, could make this configurable + + for agent in self.agents: + try: + # Add local approval OTP + addLocalApproval(api, batch, duration, agent.agentid) + # Move to audit mode + result = moveAgentToRelatedPolicy(api, agent, "audit") + successful.append((agent, result)) + logger.info( + f"Successfully moved {agent.hostname} to local approval" + ) + except Exception as e: + unsuccessful.append((agent, str(e))) + logger.error( + f"Failed to move {agent.hostname} to local approval: {e}" + ) + + except Exception as e: + logger.error(f"Error during local approval operation: {e}") + status_label.update(f"ҝŒ Error: {str(e)}") + self.operation_in_progress = False + return + + self.operation_in_progress = False + status_label.update("Òœ… Operation complete!") + + # Display results in the widget + self._display_results("Local Approval Mode", successful, unsuccessful) + + # Also post message for potential parent handling + self.post_message( + self.OperationComplete( + "Local Approval Mode", self.agents, successful, unsuccessful + ) + ) + + def _start_toggle_enforcement_operation(self) -> None: + """ + Toggle agents between enforcement and audit policy modes. + + This operation intelligently switches each agent between enforcement and + audit modes based on its current state: + - If agent.groupid is in POLICY_MAP_ENF_AUD: currently enforcing Ò†' move to audit + - Otherwise: currently in audit Ò†' move to enforcement + + The operation: + - Retrieves the enforcement/audit policy relationship map from protected config + - Sets operation state flags and updates status label + - Iterates through agents, determining current mode and toggling to opposite + - Tracks successful toggles with the new mode in the result message + - Logs both successes and failures + - Displays results and posts OperationComplete message + + The policy relationship map (POLICY_MAP_ENF_AUD) must be present in protected + configuration and maps enforcement policy IDs to audit policy IDs. If the map + is empty or not found, all agents are assumed to be in audit mode and will + be moved to enforcement. + + Returns to a non-busy state after completion regardless of individual results. + """ + self.selected_operation = "toggle_enforcement" + self.operation_in_progress = True + + status_label = self.query_one("#status_label", Static) + status_label.update("ҏ³ Toggling enforcement mode...") + + # Get API from app + api = self.app.api + + successful = [] + unsuccessful = [] + + try: + from services.agenthandler import moveAgentToRelatedPolicy + from utils.configmanager import get_protected_json + + policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD", "{}") + + for agent in self.agents: + try: + # Determine current mode and toggle + if agent.groupid in policy_relationship_map: + # Currently in enforcement, move to audit + result = moveAgentToRelatedPolicy(api, agent, "audit") + mode = "audit" + else: + # Currently in audit, move to enforcement + result = moveAgentToRelatedPolicy(api, agent, "enforcement") + mode = "enforcement" + + successful.append((agent, f"Moved to {mode}: {result}")) + logger.info(f"Successfully toggled {agent.hostname} to {mode}") + except Exception as e: + unsuccessful.append((agent, str(e))) + logger.error(f"Failed to toggle {agent.hostname}: {e}") + + except Exception as e: + logger.error(f"Error during toggle enforcement operation: {e}") + status_label.update(f"ҝŒ Error: {str(e)}") + self.operation_in_progress = False + return + + self.operation_in_progress = False + status_label.update("Òœ… Operation complete!") + + # Display results in the widget + self._display_results("Toggle Audit/Enforcement", successful, unsuccessful) + + # Also post message for potential parent handling + self.post_message( + self.OperationComplete( + "Toggle Audit/Enforcement", self.agents, successful, unsuccessful + ) + ) + + def _start_other_policy_operation(self) -> None: + """ + Move agents to a user-selected policy (currently unimplemented). + + This operation is intended to allow bulk movement of selected agents to any + alternative policy via a policy selection dialog. Currently, this feature + is not fully implemented. + + Planned Implementation: + 1. Push a new policy selector screen (TUI modal/overlay) + 2. Allow user to choose target policy from available options + 3. Move all selected agents to the chosen policy + 4. Display results like other operations + + Current Behavior: + - Sets selected_operation to "other_policy" + - Displays "Policy selection not yet implemented" status message + - Clears selected_operation without performing any action + + TODO: Complete implementation by: + - Creating a policy selector screen component + - Implementing the policy selection logic + - Integrating with moveAgentToPolicy API call + - Adding proper result tracking and display + """ + self.selected_operation = "other_policy" + self.operation_in_progress = True + + status_label = self.query_one("#status_label", Static) + status_label.update("Loading available policies...") + + try: + # Fetch all policies from API + api = self.app.api + + # Fetch all available policies + all_policies_df = api.policy_find_all() + + if all_policies_df.empty: + status_label.update("No policies available") + self.operation_in_progress = False + self.selected_operation = "" + return + + # Create and push the policy selector screen + policy_selector_screen = PolicySelectorScreen( + policies=all_policies_df, + agent_move_operations=self, + ) + self.app.push_screen(policy_selector_screen) + + except Exception as e: + logger.error(f"Error loading policies: {e}") + status_label.update(f"Error: {str(e)}") + self.operation_in_progress = False + self.selected_operation = "" + self.app.notify(f"Failed to load policies: {str(e)}", severity="error") + + def _execute_move_to_policy(self, target_policy) -> None: + """ + Execute the actual move of agents to the selected policy. + + Moves each agent sequentially to the target policy, tracking success/failure. + Updates the status label and displays results upon completion. + + Args: + target_policy: The Policy object selected by the user. + """ + status_label = self.query_one("#status_label", Static) + status_label.update(f"Moving agents to {target_policy.name}...") + + api = self.app.api + successful = [] + unsuccessful = [] + + try: + for agent in self.agents: + try: + # Move agent to target policy + result = api.agent_move(agent.agentid, target_policy.groupid) + successful.append((agent, f"Moved to {target_policy.name}")) + logger.info( + f"Successfully moved {agent.hostname} to policy {target_policy.name}" + ) + except Exception as e: + unsuccessful.append((agent, str(e))) + logger.error( + f"Failed to move {agent.hostname} to policy {target_policy.name}: {e}" + ) + + except Exception as e: + logger.error(f"Error during move to policy operation: {e}") + status_label.update(f"Error: {str(e)}") + self.operation_in_progress = False + return + + self.operation_in_progress = False + status_label.update("Operation complete!") + + # Display results in the widget + self._display_results( + f"Move to {target_policy.name}", + successful, + unsuccessful, + ) + + # Also post message for potential parent handling + self.post_message( + self.OperationComplete( + f"Move to {target_policy.name}", + self.agents, + successful, + unsuccessful, + ) + ) diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index e46c636..dd0fe1b 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -6,7 +6,15 @@ from textual.containers import Horizontal, Vertical from textual.css.query import NoMatches from textual.message import Message from textual.widget import Widget -from textual.widgets import Button, SelectionList, Static, Switch, TextArea +from textual.widgets import ( + Button, + Footer, + Header, + SelectionList, + Static, + Switch, + TextArea, +) from models.agent import Agent @@ -31,6 +39,7 @@ class MultiAgentSelector(Widget): self._match_type = value def compose(self): + yield Header(show_clock=True, icon="βš™") title_text = Static("πŸ–§ Multi-Agent Selector", id="selector_title") title_text.styles.margin = (0, 0, 0, 1) yield title_text @@ -98,6 +107,7 @@ class MultiAgentSelector(Widget): right_pane.styles.width = "2fr" yield SelectionList(id="match_results") yield Static(id="unmatched_label") + yield Footer() def on_switch_changed(self, event: Switch.Changed): self.match_type = "fuzzy" if event.value else "exact" diff --git a/widgets/policyselector.py b/widgets/policyselector.py new file mode 100644 index 0000000..1796968 --- /dev/null +++ b/widgets/policyselector.py @@ -0,0 +1,505 @@ +""" +Policy Selector Widget Module + +Provides a Textual widget for selecting target policies for bulk agent operations. +Allows users to browse available policies and select one as the destination for +moving agents. Automatically excludes parent/logical policies. +""" + +import logging +import re +from typing import Optional + +import pandas as pd +from textual.containers import Horizontal, Vertical +from textual.message import Message +from textual.widget import Widget +from textual.widgets import Button, DataTable, Footer, Header, Static, TextArea + +from models.policy import Policy + +logger = logging.getLogger(__name__) + + +class PolicySelector(Widget): + """ + A Textual widget for selecting a target policy for agent operations. + + This widget displays available policies in a table and allows users to select + one policy as the destination for bulk agent movements. It automatically excludes: + - Parent/logical policies (where parent == "global-policy-settings") + - Specified policy IDs (e.g., the current policy) + + Features: + - Wildcard filtering (* and ?) + - Interactive table for policy browsing + - Explicit confirm button for selection + - Cancel/back button to dismiss + + Attributes: + policies (list[Policy]): List of available Policy objects to display. + excluded_policy_ids (set[str]): Set of policy IDs to exclude from selection. + selected_policy (Optional[Policy]): The currently selected policy (if any). + + Automatically Filtered Out: + - Policies with parent == "global-policy-settings" (parent policies for organization) + - Any policies in excluded_policy_ids set + + Example: + ```python + policies = [policy1, policy2, policy3] + widget = PolicySelector(policies, excluded_policy_ids={current_policy.groupid}) + ``` + """ + + class PolicySelected(Message): + """ + Message posted when a policy is selected. + + Attributes: + policy (Policy): The selected policy object. + """ + + def __init__(self, policy: Policy): + super().__init__() + self.policy = policy + + def __init__(self, policies: list): + """ + Initialize the PolicySelector widget. + + Args: + policies (list): List of Policy objects or DataFrame rows to display. + Can be a list of Policy objects or a pandas DataFrame of policy data. + """ + super().__init__() + self.policies = policies + self.selected_policy: Optional[Policy] = None + self._filtered_policies = [] + self._displayed_policies = [] # Track what's currently shown in the table + self._filter_text = "" + + def compose(self): + """ + Build the UI layout for the PolicySelector widget. + + The layout includes: + - Title indicating policy selection + - Search/filter text area with wildcard support + - Filter help text showing wildcard options + - Apply Filter button + - Clear Filter button + - Confirm Selection button + - Policy table displaying available policies + - Back and Continue buttons for navigation + """ + yield Header(show_clock=True, icon="βš™") + title_text = Static( + "🎯 Select Target Policy", + id="policy_selector_title", + ) + title_text.styles.margin = (0, 0, 1, 0) + yield title_text + + with Horizontal() as main_layout: + main_layout.styles.height = "auto" + + # Left side - Filter and controls + with Vertical() as left_side: + left_side.styles.width = "1fr" + left_side.styles.height = "auto" + + filter_label = Static("Filter Policies:") + filter_label.styles.margin = (0, 0, 0, 0) + yield filter_label + + filter_input = TextArea( + id="policy_filter", + text="", + ) + filter_input.styles.height = 3 + filter_input.styles.margin = (0, 0, 1, 0) + yield filter_input + + filter_help = Static("(Use * and ? for wildcards)", id="filter_help") + filter_help.styles.margin = (0, 0, 1, 0) + yield filter_help + + apply_button = Button("βœ“ Apply Filter", id="filter_button") + apply_button.styles.width = "100%" + apply_button.styles.margin = (0, 0, 1, 0) + yield apply_button + + clear_button = Button("πŸ—‘οΈ Clear Filter", id="clear_filter_button") + clear_button.styles.width = "100%" + clear_button.styles.margin = (0, 0, 1, 0) + yield clear_button + + confirm_button = Button("βœ… Confirm Selection", id="confirm_button") + confirm_button.styles.width = "100%" + confirm_button.styles.margin = (1, 0, 1, 0) + yield confirm_button + + selected_label = Static("", id="selected_policy_label") + selected_label.styles.margin = (2, 0, 1, 0) + yield selected_label + + # Right side - Policy table + with Vertical() as right_side: + right_side.styles.width = "2fr" + right_side.styles.height = "auto" + + table_label = Static("Available Policies:") + table_label.styles.margin = (0, 0, 0, 0) + yield table_label + + policy_table = DataTable(id="policy_table", cursor_type="row") + policy_table.styles.height = "1fr" + policy_table.styles.margin = (1, 0, 1, 0) + yield policy_table + + # Bottom buttons + with Horizontal() as button_row: + button_row.styles.height = "auto" + button_row.styles.margin = (1, 0, 0, 0) + + cancel_button = Button("βœ• Cancel", id="back_button", variant="error") + cancel_button.styles.width = "1fr" + yield cancel_button + + continue_button = Button( + "β–Ά Continue", + id="continue_button", + variant="primary", + ) + continue_button.styles.width = "1fr" + continue_button.styles.margin = (0, 0, 0, 1) + yield continue_button + yield Footer() + + def on_mount(self) -> None: + """ + Initialize the policy table when the widget is mounted. + + Populates the table with column (Policy Name) and rows for each + available policy (excluding those in excluded_policy_ids and parent policies). + Sets up event handlers for table row selection. + + Filters out: + - Parent policies (where parent == "global-policy-settings") + """ + table = self.query_one("#policy_table", DataTable) + + # Configure table for row selection + table.cursor_type = "row" + table.zebra_stripes = True + + # Only add Policy Name column + table.add_columns("Policy Name") + + # Filter out excluded policies and convert to list if DataFrame + if isinstance(self.policies, pd.DataFrame): + policies_list = self.policies.to_dict("records") + else: + policies_list = self.policies + + self._filtered_policies = [] + self._displayed_policies = [] # Initialize displayed list + + for policy_data in policies_list: + # Handle both Policy objects and dict/DataFrame rows + if isinstance(policy_data, Policy): + policy_id = policy_data.groupid + policy_name = policy_data.name + parent = policy_data.parent + else: + policy_id = policy_data.get("groupid", "Unknown") + policy_name = policy_data.get("name", "Unknown") + parent = policy_data.get("parent", None) + + # Skip parent policies (logical policies that shouldn't have devices) + if parent == "global-policy-settings": + logger.debug(f"Skipping parent policy: {policy_name}") + continue + + self._filtered_policies.append(policy_data) + self._displayed_policies.append(policy_data) # Add to displayed list + + table.add_row( + policy_name, + key=policy_id, + ) + + def on_button_pressed(self, event: Button.Pressed): + """ + Handle button press events from the widget. + + Routes to: + - back_button (Cancel): Pop screen without selecting + - filter_button (Apply Filter): Filter policies with wildcard support + - clear_filter_button: Clear filter and show all policies + - confirm_button: Confirm selection and post message + - continue_button: Continue without posting message + + Args: + event (Button.Pressed): The button press event. + """ + btn_id = event.button.id + + if btn_id == "back_button": + self.app.pop_screen() + event.stop() + + elif btn_id == "filter_button": + self._apply_filter() + event.stop() + + elif btn_id == "clear_filter_button": + self._clear_filter() + event.stop() + + elif btn_id == "confirm_button": + self._confirm_selection() + event.stop() + + elif btn_id == "continue_button": + self.app.pop_screen() + event.stop() + + def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: + """ + Handle row selection in the policy table. + + Updates the selected_policy and displays the selection in the UI. + + Args: + event: DataTable.RowSelected event containing the selected row data. + """ + try: + # Get the row key from the event + row_key = event.row_key + if row_key is None: + return + + # Find the policy with matching groupid + for policy_data in self._displayed_policies: + if isinstance(policy_data, Policy): + if policy_data.groupid == row_key.value: + self.selected_policy = policy_data + break + else: + if policy_data.get("groupid") == row_key.value: + self.selected_policy = Policy( + groupid=policy_data.get("groupid"), + hidden=policy_data.get("hidden", False), + name=policy_data.get("name"), + parent=policy_data.get("parent"), + ) + break + + if self.selected_policy: + # Update selection display + label = self.query_one("#selected_policy_label", Static) + label.update(f"βœ“ Selected: {self.selected_policy.name}") + + # Log for debugging + logger.debug( + f"Selected policy: {self.selected_policy.name} (ID: {self.selected_policy.groupid})" + ) + self.app.notify( + f"Selected: {self.selected_policy.name}", + severity="information", + timeout=1, + ) + + except Exception as e: + logger.error(f"Error handling row selection: {e}") + self.app.notify(f"Selection error: {str(e)}", severity="error") + + def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None: + """ + Handle row highlighting (cursor movement) in the table. + + This provides immediate visual feedback when navigating rows. + """ + try: + # Get the row key from the event + row_key = event.row_key + if row_key is None: + return + + # Find the highlighted policy + highlighted_name = None + for policy_data in self._displayed_policies: + if isinstance(policy_data, Policy): + if policy_data.groupid == row_key.value: + highlighted_name = policy_data.name + break + else: + if policy_data.get("groupid") == row_key.value: + highlighted_name = policy_data.get("name") + break + + if highlighted_name: + label = self.query_one("#selected_policy_label", Static) + label.update(f"β†’ Highlighting: {highlighted_name}") + + except Exception as e: + logger.error(f"Error handling row highlight: {e}") + + def _apply_filter(self) -> None: + """ + Apply filter text to policy list with wildcard support. + + Supports wildcards: + - * matches any sequence of characters + - ? matches a single character + + Examples: + - "policy*" matches "policy_prod", "policy_dev", etc. + - "policy?" matches "policy1", "policy2", etc. + - "*audit*" matches anything containing "audit" + - "*test*" matches "AT Testing", "test_policy", etc. + + Filters policies by name or ID (case-insensitive) and refreshes the table display + with only matching policies. Only filters from already-filtered list + (which excludes parent policies and excluded IDs). + """ + try: + filter_input = self.query_one("#policy_filter", TextArea) + filter_text = filter_input.text.strip() + + table = self.query_one("#policy_table", DataTable) + table.clear() + + # Clear the displayed policies list + self._displayed_policies = [] + + # Compile wildcard pattern if filter text is provided + pattern = None + if filter_text: + # Escape special regex chars but preserve wildcards + pattern_text = re.escape(filter_text.lower()) + pattern_text = pattern_text.replace(r"\*", ".*").replace(r"\?", ".") + # Use search() for partial matching + pattern = re.compile(pattern_text, re.IGNORECASE) + + # Filter policies based on search text + for policy_data in self._filtered_policies: + # Handle both Policy objects and dict/DataFrame rows + if isinstance(policy_data, Policy): + policy_name = policy_data.name.lower() + policy_id = policy_data.groupid.lower() + display_name = policy_data.name + key_id = policy_data.groupid + else: + policy_name = str(policy_data.get("name", "")).lower() + policy_id = str(policy_data.get("groupid", "Unknown")).lower() + display_name = policy_data.get("name") + key_id = policy_data.get("groupid") + + # Match against filter text with wildcard support + if pattern: + # Use search() for partial matching + matches = pattern.search(policy_name) or pattern.search(policy_id) + else: + matches = True + + if matches: + # Add to displayed policies list + self._displayed_policies.append(policy_data) + + # Add row to table + table.add_row( + display_name, + key=key_id, + ) + + displayed_count = len(self._displayed_policies) + status_text = f"πŸ“Š Showing {displayed_count} of {len(self._filtered_policies)} policies" + self.app.notify(status_text, severity="information", timeout=2) + + # Clear selection when filter is applied + self.selected_policy = None + label = self.query_one("#selected_policy_label", Static) + label.update("") + + except Exception as e: + logger.error(f"Error applying filter: {e}") + self.app.notify(f"❌ Filter error: {str(e)}", severity="error") + + def _clear_filter(self) -> None: + """ + Clear the filter and display all available policies. + + Resets the filter text and refreshes the table to show all policies + (already excluding parent policies and excluded IDs). + """ + try: + filter_input = self.query_one("#policy_filter", TextArea) + filter_input.text = "" + + table = self.query_one("#policy_table", DataTable) + table.clear() + + # Reset displayed policies to all filtered policies + self._displayed_policies = list(self._filtered_policies) + + # Reload all policies + for policy_data in self._filtered_policies: + if isinstance(policy_data, Policy): + policy_id = policy_data.groupid + policy_name = policy_data.name + else: + policy_id = policy_data.get("groupid", "Unknown") + policy_name = policy_data.get("name", "Unknown") + + # Add row with only policy name + table.add_row( + policy_name, + key=policy_id, + ) + + self.selected_policy = None + label = self.query_one("#selected_policy_label", Static) + label.update("") + + except Exception as e: + logger.error(f"Error clearing filter: {e}") + + def on_text_area_changed(self, event) -> None: + """ + Handle TextArea change events - specifically for Enter key in filter. + + When the user types in the filter TextArea and the text ends with a newline, + treat it as pressing Enter and apply the filter. + """ + if event.text_area.id == "policy_filter": + # Check if the text ends with a newline (Enter was pressed) + if event.text_area.text.endswith("\n"): + # Remove the newline that was added + event.text_area.text = event.text_area.text.rstrip("\n") + # Apply the filter + self._apply_filter() + + def _confirm_selection(self) -> None: + """ + Confirm the selected policy and post selection message. + + Posts a PolicySelected message to the parent widget/screen with the + selected policy. If no policy is selected, displays an error notification. + """ + if self.selected_policy is None: + self.app.notify( + "⚠️ Please select a policy first by clicking on a row in the table", + severity="warning", + timeout=3, + ) + return + + # Log confirmation for debugging + logger.info(f"Confirming selection of policy: {self.selected_policy.name}") + self.app.notify( + f"βœ… Confirmed: {self.selected_policy.name}", severity="success", timeout=2 + ) + self.post_message(self.PolicySelected(self.selected_policy)) diff --git a/widgets/resultsdisplay.py b/widgets/resultsdisplay.py new file mode 100644 index 0000000..ebabce7 --- /dev/null +++ b/widgets/resultsdisplay.py @@ -0,0 +1,178 @@ +import logging + +from textual.containers import Horizontal, Vertical +from textual.message import Message +from textual.widget import Widget +from textual.widgets import Button, Footer, Header, Static + +logger = logging.getLogger(__name__) + + +class ResultsDisplay(Widget): + """Widget for displaying operation results in a two-column layout.""" + + CSS = """ + ResultsDisplay { + height: 100%; + } + + #results_screen { + height: 100%; + } + + #results_title { + text-align: center; + margin: 1 0; + text-style: bold; + } + + #results_layout { + height: 1fr; + margin: 1 0; + } + + #left_column, #right_column { + width: 1fr; + height: 100%; + border: solid green; + padding: 1; + } + + #right_column { + border: solid red; + } + + #success_label, #failure_label { + text-style: bold; + margin-bottom: 1; + } + + #success_results, #failure_results { + height: 1fr; + overflow-y: auto; + background: $surface; + border: round $primary; + padding: 1; + } + + .copy_button { + margin-top: 1; + width: 100%; + } + + #button_row { + height: auto; + margin: 1 0 0 0; + } + + #back_button { + width: 1fr; + } + """ + + class CopySuccess(Message): + """Posted when success results are copied.""" + + pass + + class CopyFailure(Message): + """Posted when failure results are copied.""" + + pass + + class GoBack(Message): + """Posted when back button is pressed.""" + + pass + + def __init__( + self, operation: str, successful_results: str, unsuccessful_results: str + ) -> None: + super().__init__() + self.operation = operation + self.successful_results = successful_results + self.unsuccessful_results = unsuccessful_results + + def compose(self): + with Vertical(id="results_screen"): + yield Header(show_clock=True, icon="βš™") + # Title + title = Static(f"πŸ“Š {self.operation} - Results", id="results_title") + yield title + + # Two-column layout + with Horizontal(id="results_layout"): + # Left Column - Success + with Vertical(id="left_column"): + yield Static("βœ… Successful", id="success_label") + yield Static(self.successful_results, id="success_results") + yield Button( + "πŸ“‹βœ… Copy Success List", + id="copy_success", + classes="copy_button", + ) + + # Right Column - Failure + with Vertical(id="right_column"): + yield Static("❌ Failed", id="failure_label") + yield Static(self.unsuccessful_results, id="failure_results") + yield Button( + "πŸ“‹βŒ Copy Failure List", + id="copy_failure", + classes="copy_button", + ) + + # Back Button + with Horizontal(id="button_row"): + back_button = Button("← Back", id="back_button") + yield back_button + yield Footer() + + def on_button_pressed(self, event: Button.Pressed) -> None: + btn_id = event.button.id + + if btn_id == "copy_success": + success_widget = self.query_one("#success_results", Static) + try: + import pyperclip + + pyperclip.copy(str(success_widget.renderable)) + self.app.notify( + "Òœ… Success list copied to clipboard!", + severity="information", + timeout=2, + ) + self.post_message(self.CopySuccess()) + except ImportError: + self.app.notify( + "Òő ï¸ pyperclip not installed. Run: pip install pyperclip", + severity="warning", + ) + except Exception as e: + self.app.notify(f"ҝŒ Failed to copy: {str(e)}", severity="error") + event.stop() + + elif btn_id == "copy_failure": + failure_widget = self.query_one("#failure_results", Static) + try: + import pyperclip + + pyperclip.copy(str(failure_widget.renderable)) + self.app.notify( + "Òœ… Failure list copied to clipboard!", + severity="information", + timeout=2, + ) + self.post_message(self.CopyFailure()) + except ImportError: + self.app.notify( + "Òő ï¸ pyperclip not installed. Run: pip install pyperclip", + severity="warning", + ) + except Exception as e: + self.app.notify(f"ҝŒ Failed to copy: {str(e)}", severity="error") + event.stop() + + elif btn_id == "back_button": + self.app.pop_screen() + event.stop() diff --git a/widgets/retro_terminal_theme.py b/widgets/retro_terminal_theme.py new file mode 100644 index 0000000..3ff0d02 --- /dev/null +++ b/widgets/retro_terminal_theme.py @@ -0,0 +1,38 @@ +from textual.color import Color + + +def get_retro_terminal_theme(): + from textual.theme import Theme + + return Theme( + name="retro-terminal", + background=Color.parse("#000000"), + primary=Color.parse("#00ff00"), + secondary=Color.parse("#00aa00"), + success=Color.parse("#00ff00"), + warning=Color.parse("#ffff00"), + error=Color.parse("#ff0000"), + surface=Color.parse("#111111"), + ) + + +RETRO_TERMINAL_CSS = """ +/* Retro terminal CRT effect */ +Screen { + align: center middle; + background: $background; + color: $text; +} + +/* Blocky, pixelated widgets */ +.widget { + border: tall $primary; + background: $surface; + width: 80%; +} + +/* Monospaced font */ +* { + font-family: "Courier New", monospace; +} +""" diff --git a/widgets/themeselector.py b/widgets/themeselector.py index fda7cd1..6a423c2 100644 --- a/widgets/themeselector.py +++ b/widgets/themeselector.py @@ -15,26 +15,25 @@ class ThemeSelector(Widget): self.theme_name = theme_name AVAILABLE_THEMES = [ - ("textual-dark", "textual-dark"), - ("textual-light", "textual-light"), - ("nord", "nord"), - ("gruvbox", "gruvbox"), - ("catppuccin-mocha", "catppuccin-mocha"), - ("dracula", "dracula"), - ("tokyo-night", "tokyo-night"), - ("monokai", "monokai"), - ("flexoki", "flexoki"), - ("catppuccin-latte", "catppuccin-latte"), - ("solarized-light", "solarized-light"), + ("Textual Dark", "textual-dark"), + ("Textual Light", "textual-light"), + ("Nord", "nord"), + ("Gruvbox", "gruvbox"), + ("Catppuccin Mocha", "catppuccin-mocha"), + ("Dracula", "dracula"), + ("Tokyo Night", "tokyo-night"), + ("Monokai", "monokai"), + ("Flexoki", "flexoki"), + ("Catppuccin Latte", "catppuccin-latte"), + ("Solarized Light", "solarized-light"), + ("Retro Terminal", "retro-terminal"), # your custom theme ] def compose(self): yield Static("Theme Options", id="theme_title") - with Vertical() as column: column.styles.width = "1fr" column.styles.height = "auto" - for label, btn_id in self.AVAILABLE_THEMES: yield Button(label, id=f"set_theme_{btn_id}", compact=True) From d9cbce617561f5538caa6fe5de6e399b04bb11b0 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Mon, 10 Nov 2025 15:17:07 -0500 Subject: [PATCH 2/4] Continuing work removing legacy functions and implementing UI changes --- airlock_libs/airlock_libs.pyi | 2 +- screens/otpworkflowscreen.py | 33 ++---- utils/tui.py | 51 +++------ widgets/OTP_generate.py | 5 +- widgets/agentmoveoperations.py | 182 ++++++++++++++++++++------------ widgets/amber_terminal_theme.py | 35 ++++++ widgets/multiagentselector.py | 2 +- widgets/policyselector.py | 32 ++---- widgets/retro_terminal_theme.py | 2 +- widgets/themeselector.py | 3 +- 10 files changed, 188 insertions(+), 159 deletions(-) create mode 100644 widgets/amber_terminal_theme.py diff --git a/airlock_libs/airlock_libs.pyi b/airlock_libs/airlock_libs.pyi index 84bee85..137dc92 100644 --- a/airlock_libs/airlock_libs.pyi +++ b/airlock_libs/airlock_libs.pyi @@ -30,7 +30,7 @@ def history_logging( checkpoint_number: str, policy_names: str, ) -> List[Dict[str, Any]]: - """ + """ Query execution history logs from the Airlock API. Parameters diff --git a/screens/otpworkflowscreen.py b/screens/otpworkflowscreen.py index fa5afc1..f3710b3 100644 --- a/screens/otpworkflowscreen.py +++ b/screens/otpworkflowscreen.py @@ -1,41 +1,24 @@ +# otp_workflow_screen.py + from typing import List from textual.app import ComposeResult from textual.screen import Screen from models.agent import Agent -from widgets.multiagentselector import MultiAgentSelector from widgets.OTP_generate import OTPGenerator class OTPWorkflowScreen(Screen): - """Screen that handles the OTP generation workflow.""" + """Screen that handles the OTP generation workflow without agent selection.""" - def __init__(self, all_agents: List[Agent]): + def __init__(self, selected_agents: List[Agent]): super().__init__() - self.all_agents = all_agents - self.selected_devices = None + self.selected_agents = selected_agents def compose(self) -> ComposeResult: - """Start with the multi-agent selector.""" - yield MultiAgentSelector(self.all_agents) - - def on_multi_agent_selector_agents_selected( - self, message: MultiAgentSelector.AgentsSelected - ) -> None: - """Handle selected agents - switch to OTP generator.""" - self.selected_devices = message.selected_agents - - # Remove the MultiAgentSelector - selector = self.query_one(MultiAgentSelector) - selector.remove() - - # Mount the OTPGenerator with the selected Agent objects - # No need to pass API - it will access self.app.api directly - self.mount(OTPGenerator(self.selected_devices)) + """Directly show the OTP generator for the selected agents.""" + yield OTPGenerator(self.selected_agents) def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: - """Handle OTP generation request - call the actual OTP generation function.""" - # This will be handled by the main app, but we can also do it here - # For now, just pass it up to the app level - pass + """Handle OTP generation request - pass it up to the app level if needed.""" diff --git a/utils/tui.py b/utils/tui.py index 9fb1c54..a930d03 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -20,18 +20,17 @@ from textual.widgets import ( from flows.otp import otp_activities_by_agent, otp_revoke from flows.prepPolicy import menu_policy_enforce -from flows.quietAgent import findQuietAgents from models.agent import Agent from models.policy import Policy from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen from screens.otpworkflowscreen import OTPWorkflowScreen -from services.agenthandler import findAgents, moveAgents, toggleEnforcement from services.API import AirlockAPIWrapper from services.policyhandler import confirmUpdateAfromE from utils.configmanager import load_env from utils.setup import get_base_directory, load_user_config from utils.utils import open_directory from widgets.agentmoveoperations import AgentMoveOperations +from widgets.amber_terminal_theme import get_amber_terminal_theme from widgets.multiagentselector import MultiAgentSelector from widgets.OTP_generate import OTPGenerator from widgets.policytreewidget import PolicyTreeWidget @@ -105,24 +104,18 @@ class MainMenuScreen(Screen): current_tab = reactive("") BUTTON_DEFS = { - "find": [ - ("πŸ” - Device Search", "find_device_button"), + "agent_actions": [ + ( + "πŸ–₯️ - Find, Move, or Generate OTP for Agents", + "move_agent_workflow_button", + ), ("πŸ”‡ - Find Quiet Hosts", "find_quiet_button"), ], - "move": [ - ("πŸ”„ - Move Agent Workflow", "move_agent_workflow_button"), - ("βœ… - Move to local approval", "move_local_button"), - ("πŸ”„ - Move to Audit/Enforcement", "move_audit_button"), - ("πŸ”€ - Move - Other", "move_other_button"), - ], - "otp": [ - ("🎫 - Generate OTPs", "otp_generate_button"), - ("πŸ“Š - OTP Activities By Agent", "otp_activities_button"), - ("❌ - Revoke OTPs", "otp_revoke_button"), - ], "policy": [ ("πŸ”’ - Prepare Policy For Enforcement", "policy_prep_button"), ("πŸ”„ - Update Audit Policies", "policy_audit_update_button"), + ("πŸ“Š - OTP Activities By Agent", "otp_activities_button"), + ("❌ - Revoke OTPs", "otp_revoke_button"), ], } @@ -148,23 +141,21 @@ class MainMenuScreen(Screen): yield Header(show_clock=True, icon="βš™") tabs = [ - Tab("Policy Tree", id="p_tree"), - Tab("Device Search", id="find"), - Tab("Move Agent", id="move"), - Tab("OTP", id="otp"), + Tab("Tree View", id="p_tree"), + Tab("Agents", id="agent_actions"), Tab("Directory", id="dir"), Tab("Settings", id="settings"), ] if self.extras == "POLICYPREP": - tabs.insert(3, Tab("Policy Prep", id="policy")) + tabs.insert(2, Tab("Policy Prep", id="policy")) yield Tabs(*tabs, id="tabs") yield Vertical(id="content") yield Footer() def on_mount(self) -> None: - self.switch_tab("find") + self.switch_tab("agent_actions") # focus helpers def _get_content_buttons(self) -> list[Button]: @@ -225,7 +216,7 @@ class MainMenuScreen(Screen): def on_multi_agent_selector_agents_selected( self, message: MultiAgentSelector.AgentsSelected ) -> None: - """Handle selected agents from MultiAgentSelector.""" + """Handle selected agents from AgentSelector.""" global _PENDING_JOB selected_agents = message.selected_agents logger.info("Selected agents: %s", selected_agents) @@ -324,26 +315,11 @@ class MainMenuScreen(Screen): logger.debug("Button pressed: %s", button_id) match button_id: - case "find_device_button": - _PENDING_JOB = ("legacy", findAgents, (self.app.api, False), {}) - case "find_quiet_button": - _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) case "move_agent_workflow_button": # Push Move Agent workflow screen self.app.push_screen(MoveAgentWorkflowScreen(self.app.devices)) event.stop() return # Don't exit the app - case "move_local_button": - _PENDING_JOB = ( - "legacy", - print, - ("Move to local approval (placeholder)",), - {}, - ) - case "move_audit_button": - _PENDING_JOB = ("legacy", toggleEnforcement, (self.app.api,), {}) - case "move_other_button": - _PENDING_JOB = ("legacy", moveAgents, (self.app.api,), {}) case "otp_generate_button": # NEW: Push OTP workflow screen instead of legacy function self.app.push_screen(OTPWorkflowScreen(self.app.devices)) @@ -415,6 +391,7 @@ class Loxide(App): def on_mount(self, api: AirlockAPIWrapper) -> None: self.register_theme(get_retro_terminal_theme()) + self.register_theme(get_amber_terminal_theme()) self.theme = self._textual_theme self.push_screen(MainMenuScreen(api)) diff --git a/widgets/OTP_generate.py b/widgets/OTP_generate.py index a3476c6..221394b 100644 --- a/widgets/OTP_generate.py +++ b/widgets/OTP_generate.py @@ -194,7 +194,10 @@ class OTPGenerator(Widget): btn_id = event.button.id if btn_id == "back_button": - self.app.pop_screen() + + while len(self.app.screen_stack) > 2: + self.app.pop_screen() + event.stop() elif btn_id == "copy_clipboard_button": diff --git a/widgets/agentmoveoperations.py b/widgets/agentmoveoperations.py index 1e55fdf..6cfe7f3 100644 --- a/widgets/agentmoveoperations.py +++ b/widgets/agentmoveoperations.py @@ -18,9 +18,13 @@ Dependencies: - flows.localApproval: Local approval workflow handling """ +from dataclasses import asdict +from datetime import datetime import logging +import os from typing import List +import pandas as pd from textual.containers import Horizontal, Vertical from textual.css.query import NoMatches from textual.message import Message @@ -29,7 +33,9 @@ from textual.widget import Widget from textual.widgets import Button, DataTable, Header, Static, TextArea from models.agent import Agent +from screens.otpworkflowscreen import OTPWorkflowScreen from screens.policyselectorscreen import PolicySelectorScreen +from widgets.OTP_generate import OTPGenerator logger = logging.getLogger(__name__) @@ -133,18 +139,24 @@ class AgentMoveOperations(Widget): Handles NoMatches exceptions gracefully in case buttons are not yet rendered. """ try: + export_csv_btn = self.query_one("#export_csv_btn", Button) local_approval_btn = self.query_one("#local_approval_btn", Button) toggle_enforcement_btn = self.query_one("#toggle_enforcement_btn", Button) other_policy_btn = self.query_one("#other_policy_btn", Button) + otp_gen_btn = self.query_one("#otp_gen_btn", Button) # If operation in progress, disable all if self.operation_in_progress: + otp_gen_btn = True + export_csv_btn.disabled = True local_approval_btn.disabled = True toggle_enforcement_btn.disabled = True other_policy_btn.disabled = True else: - # If an operation was selected, keep it disabled, enable others + # If an operation was selected, disable if self.selected_operation: + otp_gen_btn.disabled = self.selected_operation == "otp_gen" + export_csv_btn.disabled = self.selected_operation == "export_csv" local_approval_btn.disabled = ( self.selected_operation == "local_approval" ) @@ -156,6 +168,8 @@ class AgentMoveOperations(Widget): ) else: # Enable all buttons + otp_gen_btn = False + export_csv_btn = False local_approval_btn.disabled = False toggle_enforcement_btn.disabled = False other_policy_btn.disabled = False @@ -189,21 +203,21 @@ class AgentMoveOperations(Widget): f"Operation: {operation_name}", f"{'=' * 50}", "", - f"Òœ… Successful ({len(successful)}):", + f"βœ… Successful ({len(successful)}):", ] if successful: for agent, result in successful: - results_lines.append(f" Ò€’ {agent.hostname}") + results_lines.append(f" βœ… {agent.hostname}") else: results_lines.append(" (none)") results_lines.append("") - results_lines.append(f"ҝŒ Failed ({len(unsuccessful)}):") + results_lines.append(f"❌ Failed ({len(unsuccessful)}):") if unsuccessful: for agent, error in unsuccessful: - results_lines.append(f" Ò€’ {agent.hostname}: {error}") + results_lines.append(f" ❌ {agent.hostname}: {error}") else: results_lines.append(" (none)") @@ -231,7 +245,7 @@ class AgentMoveOperations(Widget): It builds a two-column layout with: - Left side: Agent table showing selected agents and their current policies - Right side: Operation buttons and results display area - - Bottom: Navigation buttons (Back, Reset) + - Bottom: Navigation buttons (Back) The layout is responsive with: - Agent table: 2/3 width @@ -240,7 +254,7 @@ class AgentMoveOperations(Widget): """ yield Header(show_clock=True, icon="βš™") title_text = Static( - f"↔️ Move Agent Operations - {len(self.agents)} device(s) selected", + f"πŸ–₯️ Agent Operations - {len(self.agents)} device(s) selected", id="move_ops_title", ) title_text.styles.margin = (0, 0, 1, 0) @@ -251,7 +265,7 @@ class AgentMoveOperations(Widget): # Left side - Agent list with Vertical() as left_side: - left_side.styles.width = "2fr" + left_side.styles.width = "3fr" left_side.styles.height = "auto" agents_label = Static("Selected Agents:") @@ -266,7 +280,8 @@ class AgentMoveOperations(Widget): # Right side - Operation buttons with Vertical() as right_side: - right_side.styles.width = "1fr" + right_side.styles.width = "2fr" + right_side.styles.margin = (0, 1, 0, 1) right_side.styles.height = "auto" operations_label = Static("Operations:") @@ -274,13 +289,23 @@ class AgentMoveOperations(Widget): yield operations_label # Operation buttons + export_csv_btn = Button("πŸ“ˆ Export CSV", id="export_csv_btn") + export_csv_btn.styles.width = "100%" + export_csv_btn.styles.margin = (0, 0, 1, 0) + yield export_csv_btn + local_approval_btn = Button( - "βœ… Local Approval Mode", id="local_approval_btn" + "βœ”οΈ Local Approval Mode", id="local_approval_btn" ) local_approval_btn.styles.width = "100%" local_approval_btn.styles.margin = (0, 0, 1, 0) yield local_approval_btn + otp_gen_btn = Button("🎫 Generate One Time Passes", id="otp_gen_btn") + otp_gen_btn.styles.width = "100%" + otp_gen_btn.styles.margin = (0, 0, 1, 0) + yield otp_gen_btn + toggle_enforcement_btn = Button( "πŸ”„ Toggle Audit/Enforcement", id="toggle_enforcement_btn" ) @@ -300,39 +325,10 @@ class AgentMoveOperations(Widget): status_label.styles.margin = (2, 0, 0, 0) yield status_label - # Results display area (initially hidden) - with Vertical(id="results_container") as results_container: - results_container.styles.height = "auto" - results_container.styles.margin = (1, 0, 0, 0) - results_container.styles.display = "none" - - results_label = Static("Γ°ΕΈβ€œΕ  Results:", id="results_label") - results_label.styles.margin = (0, 0, 0, 0) - yield results_label - - results_text = TextArea(id="results_text", read_only=True) - results_text.styles.height = 15 - results_text.styles.margin = (0, 0, 1, 0) - yield results_text - - copy_results_btn = Button( - "Γ°ΕΈβ€œβ€Ή Copy Results to Clipboard", id="copy_results_btn" - ) - copy_results_btn.styles.width = "100%" - yield copy_results_btn - - # Bottom buttons - with Horizontal() as button_row: - button_row.styles.height = "auto" - button_row.styles.margin = (1, 0, 0, 0) - - back_button = Button("Ò† Back", id="back_button") - back_button.styles.width = "1fr" - yield back_button - - reset_button = Button("Γ°ΕΈβ€β€ž Reset Selection", id="reset_button") - reset_button.styles.width = "1fr" - yield reset_button + back_button = Button("← Back", id="back_button") + back_button.styles.width = "50%" + back_button.styles.margin = (0, 1, 1, 0) + yield back_button def on_mount(self) -> None: """ @@ -359,13 +355,15 @@ class AgentMoveOperations(Widget): self._update_button_states() + def on_otp_generator_otp_info(self, message: OTPGenerator.OTPInfo) -> None: + """Handle OTP generation request - call the actual OTP generation function.""" + def on_button_pressed(self, event: Button.Pressed): """ Handle button press events from the widget. This Textual event handler routes button presses to appropriate actions: - back_button: Pop this screen (return to parent) - - reset_button: Clear operation state and hide results - copy_results_btn: Copy results text to clipboard (requires pyperclip) - local_approval_btn: Start local approval operation - toggle_enforcement_btn: Start toggle audit/enforcement operation @@ -380,21 +378,8 @@ class AgentMoveOperations(Widget): btn_id = event.button.id if btn_id == "back_button": - self.app.pop_screen() - event.stop() - - elif btn_id == "reset_button": - # Reset operation selection - self.selected_operation = "" - self.operation_in_progress = False - status_label = self.query_one("#status_label", Static) - status_label.update("") - # Hide results - try: - results_container = self.query_one("#results_container", Vertical) - results_container.styles.display = "none" - except NoMatches: - pass + while len(self.app.screen_stack) > 2: + self.app.pop_screen() event.stop() elif btn_id == "copy_results_btn": @@ -404,18 +389,21 @@ class AgentMoveOperations(Widget): pyperclip.copy(results_text.text) self.app.notify( - "Òœ… Results copied to clipboard!", + "πŸ“‹βœ… Results copied to clipboard!", severity="information", timeout=2, ) except ImportError: self.app.notify( - "Òő ï¸ pyperclip not installed. Run: pip install pyperclip", + "❌ pyperclip not installed. Run: pip install pyperclip", severity="warning", ) except Exception as e: self.app.notify(f"ҝŒ Failed to copy: {str(e)}", severity="error") event.stop() + elif btn_id == "export_csv_btn": + self._start_export_csv_operation() + event.stop() elif btn_id == "local_approval_btn": self._start_local_approval_operation() @@ -428,6 +416,9 @@ class AgentMoveOperations(Widget): elif btn_id == "other_policy_btn": self._start_other_policy_operation() event.stop() + elif btn_id == "otp_gen_btn": + self._start_OTP_gen_operation() + event.stop() def _start_local_approval_operation(self) -> None: """ @@ -456,7 +447,7 @@ class AgentMoveOperations(Widget): self.operation_in_progress = True status_label = self.query_one("#status_label", Static) - status_label.update("ҏ³ Moving agents to local approval...") + status_label.update("βœ”οΈ Moving agents to local approval...") # Get API from app api = self.app.api @@ -491,12 +482,12 @@ class AgentMoveOperations(Widget): except Exception as e: logger.error(f"Error during local approval operation: {e}") - status_label.update(f"ҝŒ Error: {str(e)}") + status_label.update(f"❌ Error: {str(e)}") self.operation_in_progress = False return self.operation_in_progress = False - status_label.update("Òœ… Operation complete!") + status_label.update("βœ… Operation complete!") # Display results in the widget self._display_results("Local Approval Mode", successful, unsuccessful) @@ -508,14 +499,63 @@ class AgentMoveOperations(Widget): ) ) + def _start_export_csv_operation(self) -> None: + self.selected_operation = "export_csv" + self.operation_in_progress = True + successful = [] + unsuccessful = [] + status_label = self.query_one("#status_label", Static) + status_label.update("Exporting CSV...") + agents = self.agents + policies = self.app.policies + path = self.app.working_dir + + try: + # Enrich each agent with policies and status text + for agent in agents: + agent.enrich_with_policies(policies) + + # Convert each Agent to a dictionary, including all fields + data = [] + for agent in agents: + row = asdict(agent) + # Remove the class-level status_map from the row + row.pop("status_map", None) + data.append(row) + + # Create DataFrame + df = pd.DataFrame(data) + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + filename = f"agentsearch_{timestamp}.csv" + file_path = os.path.join(str(path), filename) + df.to_csv(file_path, index=False) + successful.append(file_path) + status_label.update(f"βœ… Exported to {file_path}") + except Exception: + status_label.update("❌ Failed") + + self.operation_in_progress = False + + """ + # Display results in the widget + self._display_results("CSV Export", successful, unsuccessful) + + # Also post message for potential parent handling + self.post_message( + self.OperationComplete( + "CSV Export", self.agents, successful, unsuccessful + ) + ) + """ + def _start_toggle_enforcement_operation(self) -> None: """ Toggle agents between enforcement and audit policy modes. This operation intelligently switches each agent between enforcement and audit modes based on its current state: - - If agent.groupid is in POLICY_MAP_ENF_AUD: currently enforcing Ò†' move to audit - - Otherwise: currently in audit Ò†' move to enforcement + - If agent.groupid is in POLICY_MAP_ENF_AUD: currently enforcing , move to audit + - Otherwise: currently in audit, move to enforcement The operation: - Retrieves the enforcement/audit policy relationship map from protected config @@ -570,7 +610,7 @@ class AgentMoveOperations(Widget): except Exception as e: logger.error(f"Error during toggle enforcement operation: {e}") - status_label.update(f"ҝŒ Error: {str(e)}") + status_label.update(f"❌ Error: {str(e)}") self.operation_in_progress = False return @@ -640,11 +680,17 @@ class AgentMoveOperations(Widget): except Exception as e: logger.error(f"Error loading policies: {e}") - status_label.update(f"Error: {str(e)}") + status_label.update(f"❌ Error: {str(e)}") self.operation_in_progress = False self.selected_operation = "" self.app.notify(f"Failed to load policies: {str(e)}", severity="error") + def _start_OTP_gen_operation(self) -> None: + status_label = self.query_one("#status_label", Static) + status_label.update("Generating OTP.") + + self.app.push_screen(OTPWorkflowScreen(self.agents)) + def _execute_move_to_policy(self, target_policy) -> None: """ Execute the actual move of agents to the selected policy. diff --git a/widgets/amber_terminal_theme.py b/widgets/amber_terminal_theme.py new file mode 100644 index 0000000..169c325 --- /dev/null +++ b/widgets/amber_terminal_theme.py @@ -0,0 +1,35 @@ +from textual.color import Color +from textual.theme import Theme + + +def get_amber_terminal_theme(): + """Amber CRT theme with compensated brightness for blending.""" + return Theme( + name="amber-terminal", + background=Color.parse("#000000"), # pure black + primary=Color.parse("#ffb733"), # bright amber + secondary=Color.parse("#e69500"), # strong amber + success=Color.parse("#ffb733"), + warning=Color.parse("#ffff66"), + error=Color.parse("#ff3300"), + surface=Color.parse("#3a1f00"), # brighter brown for blending + ) + + +AMBER_TERMINAL_CSS = """ +Screen { + align: center middle; + background: #000000; /* force black */ + color: #ffb733; /* force amber text */ +} + +.widget { + border: tall #ffb733; /* force amber border */ + background: #3a1f00; /* compensated surface */ + width: 80%; +} + +* { + font-family: "Courier New", monospace; +} +""" diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index dd0fe1b..a6c97e0 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -40,7 +40,7 @@ class MultiAgentSelector(Widget): def compose(self): yield Header(show_clock=True, icon="βš™") - title_text = Static("πŸ–§ Multi-Agent Selector", id="selector_title") + title_text = Static("πŸ–§ Agent Selector", id="selector_title") title_text.styles.margin = (0, 0, 0, 1) yield title_text diff --git a/widgets/policyselector.py b/widgets/policyselector.py index 1796968..c5bc6f7 100644 --- a/widgets/policyselector.py +++ b/widgets/policyselector.py @@ -91,7 +91,7 @@ class PolicySelector(Widget): - Clear Filter button - Confirm Selection button - Policy table displaying available policies - - Back and Continue buttons for navigation + - Back buttons for navigation """ yield Header(show_clock=True, icon="βš™") title_text = Static( @@ -144,6 +144,11 @@ class PolicySelector(Widget): selected_label.styles.margin = (2, 0, 1, 0) yield selected_label + cancel_button = Button("← Back", id="back_button") + cancel_button.styles.width = "100%" + cancel_button.styles.margin = (1, 0, 1, 0) + yield cancel_button + # Right side - Policy table with Vertical() as right_side: right_side.styles.width = "2fr" @@ -158,23 +163,6 @@ class PolicySelector(Widget): policy_table.styles.margin = (1, 0, 1, 0) yield policy_table - # Bottom buttons - with Horizontal() as button_row: - button_row.styles.height = "auto" - button_row.styles.margin = (1, 0, 0, 0) - - cancel_button = Button("βœ• Cancel", id="back_button", variant="error") - cancel_button.styles.width = "1fr" - yield cancel_button - - continue_button = Button( - "β–Ά Continue", - id="continue_button", - variant="primary", - ) - continue_button.styles.width = "1fr" - continue_button.styles.margin = (0, 0, 0, 1) - yield continue_button yield Footer() def on_mount(self) -> None: @@ -239,7 +227,6 @@ class PolicySelector(Widget): - filter_button (Apply Filter): Filter policies with wildcard support - clear_filter_button: Clear filter and show all policies - confirm_button: Confirm selection and post message - - continue_button: Continue without posting message Args: event (Button.Pressed): The button press event. @@ -247,7 +234,8 @@ class PolicySelector(Widget): btn_id = event.button.id if btn_id == "back_button": - self.app.pop_screen() + while len(self.app.screen_stack) > 2: + self.app.pop_screen() event.stop() elif btn_id == "filter_button": @@ -262,10 +250,6 @@ class PolicySelector(Widget): self._confirm_selection() event.stop() - elif btn_id == "continue_button": - self.app.pop_screen() - event.stop() - def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: """ Handle row selection in the policy table. diff --git a/widgets/retro_terminal_theme.py b/widgets/retro_terminal_theme.py index 3ff0d02..e0cd9c9 100644 --- a/widgets/retro_terminal_theme.py +++ b/widgets/retro_terminal_theme.py @@ -12,7 +12,7 @@ def get_retro_terminal_theme(): success=Color.parse("#00ff00"), warning=Color.parse("#ffff00"), error=Color.parse("#ff0000"), - surface=Color.parse("#111111"), + surface=Color.parse("#071802"), ) diff --git a/widgets/themeselector.py b/widgets/themeselector.py index 6a423c2..bcbbb99 100644 --- a/widgets/themeselector.py +++ b/widgets/themeselector.py @@ -26,7 +26,8 @@ class ThemeSelector(Widget): ("Flexoki", "flexoki"), ("Catppuccin Latte", "catppuccin-latte"), ("Solarized Light", "solarized-light"), - ("Retro Terminal", "retro-terminal"), # your custom theme + ("Retro Terminal", "retro-terminal"), + ("Amber Terminal", "amber-terminal"), # your custom theme ] def compose(self): From 3c2e825210377f2cbf9b61ae9c459ca2db5cb7ec Mon Sep 17 00:00:00 2001 From: Zarithas Date: Mon, 10 Nov 2025 15:47:30 -0500 Subject: [PATCH 3/4] Minor UI Tweak --- models/policy.py | 31 +++++++++++++++++-------------- widgets/multiagentselector.py | 8 ++++---- widgets/policyselector.py | 4 +++- 3 files changed, 24 insertions(+), 19 deletions(-) diff --git a/models/policy.py b/models/policy.py index dbfb357..09e16e7 100644 --- a/models/policy.py +++ b/models/policy.py @@ -13,33 +13,36 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -import json """ Policy model representing policy data and relationships. """ -class Policy: - def __init__(self, groupid, hidden, name, parent): - self.groupid = groupid - self.hidden = hidden - self.name = name - self.parent = parent +# policy.py - def __repr__(self): - # Show all current attributes, including dynamically added ones +from dataclasses import asdict, dataclass, field +import json +from typing import Optional + + +@dataclass(order=True) +class Policy: + name: str + groupid: int = field(compare=False) + hidden: bool = field(compare=False) + parent: Optional[str] = field(default=None, compare=False) + + def __repr__(self) -> str: attrs = ", ".join( f"{key}={repr(value)}" for key, value in self.__dict__.items() ) return f"" - def to_dict(self): - # Return all attributes as a dictionary - return self.__dict__ + def to_dict(self) -> dict: + return asdict(self) - def to_json(self): - # Convert to JSON string, handling non-serializable types gracefully + def to_json(self) -> str: return json.dumps(self.to_dict(), default=str) diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index a6c97e0..92016b7 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -79,14 +79,14 @@ class MultiAgentSelector(Widget): with Horizontal() as select_buttons: select_buttons.styles.margin = (0, 0, 0, 0) - select_all_button = Button("βœ… Select All", id="select_all") - select_all_button.styles.margin = (1, 1, 0, 1) - yield select_all_button - select_none_button = Button("🚫 Select None", id="select_none") select_none_button.styles.margin = (1, 0, 0, 1) yield select_none_button + select_all_button = Button("βœ… Select All", id="select_all") + select_all_button.styles.margin = (1, 1, 0, 1) + yield select_all_button + with Horizontal() as button_row: button_row.styles.height = "auto" button_row.styles.margin = (1, 0, 0, 0) diff --git a/widgets/policyselector.py b/widgets/policyselector.py index c5bc6f7..fb34218 100644 --- a/widgets/policyselector.py +++ b/widgets/policyselector.py @@ -108,6 +108,7 @@ class PolicySelector(Widget): with Vertical() as left_side: left_side.styles.width = "1fr" left_side.styles.height = "auto" + left_side.styles.margin = (0, 1, 0, 1) filter_label = Static("Filter Policies:") filter_label.styles.margin = (0, 0, 0, 0) @@ -130,7 +131,7 @@ class PolicySelector(Widget): apply_button.styles.margin = (0, 0, 1, 0) yield apply_button - clear_button = Button("πŸ—‘οΈ Clear Filter", id="clear_filter_button") + clear_button = Button("Clear Filter", id="clear_filter_button") clear_button.styles.width = "100%" clear_button.styles.margin = (0, 0, 1, 0) yield clear_button @@ -190,6 +191,7 @@ class PolicySelector(Widget): policies_list = self.policies.to_dict("records") else: policies_list = self.policies + policies_list = sorted(policies_list) self._filtered_policies = [] self._displayed_policies = [] # Initialize displayed list From 5cb079dad74f864b222b88f05dac8ab5d4ebe8f4 Mon Sep 17 00:00:00 2001 From: Zarithas Date: Mon, 10 Nov 2025 17:21:18 -0500 Subject: [PATCH 4/4] Fixed Breaking Legacy change --- utils/tui.py | 71 ++++++++++++++-------------------- widgets/agentmoveoperations.py | 5 ++- widgets/multiagentselector.py | 5 ++- 3 files changed, 36 insertions(+), 45 deletions(-) diff --git a/utils/tui.py b/utils/tui.py index a930d03..8cadf33 100644 --- a/utils/tui.py +++ b/utils/tui.py @@ -20,6 +20,7 @@ from textual.widgets import ( from flows.otp import otp_activities_by_agent, otp_revoke from flows.prepPolicy import menu_policy_enforce +from flows.quietAgent import findQuietAgents from models.agent import Agent from models.policy import Policy from screens.moveagentworkflowscreen import MoveAgentWorkflowScreen @@ -325,6 +326,8 @@ class MainMenuScreen(Screen): self.app.push_screen(OTPWorkflowScreen(self.app.devices)) event.stop() return # Don't exit the app + case "find_quiet_button": + _PENDING_JOB = ("legacy", findQuietAgents, (self.app.api,), {}) case "otp_activities_button": _PENDING_JOB = ("legacy", otp_activities_by_agent, (self.app.api,), {}) case "otp_revoke_button": @@ -353,7 +356,6 @@ class Loxide(App): text-align: center; } """ - BINDINGS = [ ("q", "quit", "Quit"), ("d", "open_dir", "Open Directory"), @@ -367,17 +369,20 @@ class Loxide(App): if not os.path.isdir(wd): wd = os.getcwd() self.working_dir = wd + # Initial data load + self.refresh_data() - # Add error handling for API calls + def refresh_data(self) -> None: + """Public method to refresh policies and devices from the API.""" try: self.policies = [ - Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows() + Policy(**row.to_dict()) + for _, row in self.api.policy_find_all().iterrows() ] self.devices = [ - Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows() + Agent(**row.to_dict()) + for _, row in self.api.agent_find_all().iterrows() ] - - # Enrich agents with policy information if self.policies and self.devices: for agent in self.devices: agent.enrich_with_policies(self.policies) @@ -401,6 +406,8 @@ class Loxide(App): self.exit() def action_open_dir(self) -> None: + # Refresh data before proceeding + self.refresh_data() screen = self.screen_stack[-1] if isinstance(screen, MainMenuScreen): if screen.current_tab != "dir": @@ -417,7 +424,6 @@ def _restore_terminal_for_legacy() -> None: sys.stdout.write("\033[?1000l\033[?1002l\033[?1003l\033[?1006l") sys.stdout.write("\033[2J\033[H") sys.stdout.flush() - if os.name == "nt": try: import ctypes @@ -434,7 +440,6 @@ def _restore_terminal_for_legacy() -> None: def _run_legacy_job(func, args, kwargs) -> None: logger.debug("Running legacy job: %s", getattr(func, "__name__", func)) _restore_terminal_for_legacy() - try: func(*args, **kwargs) finally: @@ -449,24 +454,31 @@ def _run_legacy_job(func, args, kwargs) -> None: # --------------------------------------------------------------------------- def run_Loxide(api: AirlockAPIWrapper) -> None: global _PENDING_JOB + base_dir = get_base_directory() + env_path = base_dir / ".env" + dotenv.load_dotenv(dotenv_path=env_path, override=True) - while True: - base_dir = get_base_directory() - env_path = base_dir / ".env" - dotenv.load_dotenv(dotenv_path=env_path, override=True) + max_attempts = 5 + attempts = 0 + while attempts < max_attempts: + attempts += 1 + logger.debug("Starting job loop iteration (attempt %d)", attempts) _PENDING_JOB = None app = Loxide(api) try: app.run() except SystemExit as exc: - logger.debug("Caught SystemExit from Textual: %s", exc) + if exc.code != 0: + logger.debug("Caught SystemExit from Textual: %s", exc) + raise job = _PENDING_JOB logger.debug("After app.run(), _PENDING_JOB = %r", job) if not job: + logger.debug("No job pending, exiting loop") break if job[0] == "legacy": @@ -475,49 +487,24 @@ def run_Loxide(api: AirlockAPIWrapper) -> None: continue if job[0] == "restart": - # just loop again; fresh .env was already loaded at the top + logger.debug("Restarting job loop") continue if job[0] == "multi_agent_action": - # Handle multi-agent selection logger.info("Multi-agent action with selected agents: %s", job[1]) continue - # NEW: Handle OTP workflow if job[0] == "otp_workflow": _, devices, requestor, reasoning, duration = job - # Call your OTP generation with the parameters def otp_generate_with_params(): - - print(f"\n{'='*60}") - print("OTP GENERATION") - print(f"{'='*60}") - print(f"Requestor: {requestor}") - print(f"Reasoning: {reasoning}") - print(f"Duration: {duration} minutes") - print(f"\nGenerating OTPs for {len(devices)} devices:") - print(f"{'='*60}\n") - - # Call your actual OTP generation function - # You'll need to adapt otp_generate to accept these parameters - # For now, this is a placeholder showing the structure - for device in devices: - print(f"Device: {device}") - print(f" Requestor: {requestor}") - print(f" Reason: {reasoning}") - print(f" Duration: {duration} minutes") - # TODO: Actually call your API to generate OTP - # result = api.generate_otp(device, requestor, reasoning, duration) - print() - - print(f"{'='*60}") - print("OTP Generation Complete!") - print(f"{'='*60}") + # Your OTP logic here + pass _run_legacy_job(otp_generate_with_params, (), {}) continue + logger.error("Unknown job type: %r", job) break diff --git a/widgets/agentmoveoperations.py b/widgets/agentmoveoperations.py index 6cfe7f3..03e5da4 100644 --- a/widgets/agentmoveoperations.py +++ b/widgets/agentmoveoperations.py @@ -506,6 +506,7 @@ class AgentMoveOperations(Widget): unsuccessful = [] status_label = self.query_one("#status_label", Static) status_label.update("Exporting CSV...") + self.app.refresh_data() agents = self.agents policies = self.app.policies path = self.app.working_dir @@ -604,6 +605,8 @@ class AgentMoveOperations(Widget): successful.append((agent, f"Moved to {mode}: {result}")) logger.info(f"Successfully toggled {agent.hostname} to {mode}") + self.app.refresh_data() + except Exception as e: unsuccessful.append((agent, str(e))) logger.error(f"Failed to toggle {agent.hostname}: {e}") @@ -728,7 +731,7 @@ class AgentMoveOperations(Widget): status_label.update(f"Error: {str(e)}") self.operation_in_progress = False return - + self.app.refresh_data() self.operation_in_progress = False status_label.update("Operation complete!") diff --git a/widgets/multiagentselector.py b/widgets/multiagentselector.py index 92016b7..f4ea371 100644 --- a/widgets/multiagentselector.py +++ b/widgets/multiagentselector.py @@ -80,11 +80,11 @@ class MultiAgentSelector(Widget): select_buttons.styles.margin = (0, 0, 0, 0) select_none_button = Button("🚫 Select None", id="select_none") - select_none_button.styles.margin = (1, 0, 0, 1) + select_none_button.styles.margin = (1, 1, 0, 1) yield select_none_button select_all_button = Button("βœ… Select All", id="select_all") - select_all_button.styles.margin = (1, 1, 0, 1) + select_all_button.styles.margin = (1, 0, 0, 1) yield select_all_button with Horizontal() as button_row: @@ -93,6 +93,7 @@ class MultiAgentSelector(Widget): back_button = Button("← Back", id="back_button") back_button.styles.width = "1fr" + back_button.styles.margin = (0, 0, 0, 1) yield back_button submit_button = Button(