From 3117816f935ee4e29089815ae17b4ca9d1a26ae5 Mon Sep 17 00:00:00 2001 From: Wojciech Janota Date: Tue, 29 Nov 2022 14:09:50 +0100 Subject: [PATCH] WIP --- database.db | Bin 28672 -> 28672 bytes fleetcontrol | 29 ++++ .../__pycache__/communication.cpython-310.pyc | Bin 5965 -> 7190 bytes network/communication.py | 60 ++++++- .../__pycache__/database.cpython-310.pyc | Bin 8728 -> 9867 bytes utils/database/database.py | 159 +++++++++++++----- 6 files changed, 208 insertions(+), 40 deletions(-) diff --git a/database.db b/database.db index 22f568f094d3feda1edbb2b3ca0865385254553d..5fcb2b73b60b1f18787f9b11479005f4ee010ad6 100644 GIT binary patch delta 224 zcmZp8z}WDBae_3Xz(g5mMuCk9OZYjM_&XW+kMVbI78I!EXSZfzwN{r_o_s-GMAOQ^ z(9qII&(O?5&rr|M(5NJ}xFo+kvCKfvK+h;KrywKIf)S`%6s&skLwRMOStl9zpYosF zEU2)TpGBJ4nh|6W2hboUDQ0U<5J%X`z<`(Ehk==|h=G3@-+jI!z9OL2Nqm!w>~66z zFfc?-Yz$+HV4Un{uRnR79oyu)_A31Rrx}>}yBYXb^LO(<<39~lU(G+c+us4CUV5UV N_~vK+5(-R)EC6qEJZS&` delta 182 zcmZp8z}WDBae_1>_e2?IM(&LXOZeFs`HwO1AKNULu!o<`f|1o)U37B4y!hnD@`^yI zrwsg0fl`%wE3C8gaVTR3jmAAG=~5H diff --git a/fleetcontrol b/fleetcontrol index a3ebb26..f5c4a8f 100644 --- a/fleetcontrol +++ b/fleetcontrol @@ -58,6 +58,14 @@ def add_image(image_name: str, image_file: str, image_version: str): logger.error(f"Error adding image to the database: {str(ex)}") exit(-1) +def remove_image(image_name: str, image_version: str): + try: + db = Database(config.database_file, config.server_loglevel) + obj_to_remove = db.get_image_by_name_version_string(f"{image_name}@{image_version}") + db.delete_image(obj_to_remove) + except Exception as ex: + logger.error(f"Error removing image from the database: {str(ex)}") + exit(-1) def assign_image(image_name: str, image_version: str, client_mac_address: str): try: @@ -65,6 +73,14 @@ def assign_image(image_name: str, image_version: str, client_mac_address: str): db.assign_image_to_client(client_mac_address=client_mac_address, image_name_version_combo=f"{image_name}@{image_version}") except Exception as ex: logger.error(f"Error assigning image to a client: {str(ex)}") + exit(-1) + +def detach_image(image_name: str, image_version: str, client_mac_address: str): + try: + db = Database(config.database_file, config.server_loglevel) + db.detach_image_from_client(client_mac_address=client_mac_address, image_name_version_combo=f"{image_name}@{image_version}") + except Exception as ex: + logger.error(f"Error detaching image from the client {client_mac_address}; error was {str(ex)}") parser = argparse.ArgumentParser( @@ -75,7 +91,9 @@ parser = argparse.ArgumentParser( function_mapper = { "run": run_server, "add_image": add_image, + "remove_image": remove_image, "assign_image": assign_image, + "detach_image": detach_image } parser.add_argument("command", choices=function_mapper) @@ -95,12 +113,23 @@ if "add_image" == args.command: image_file=args.image_filepath, image_version=args.image_version, ) +elif "remove_image" == args.command: + fun( + image_name=args.image_name, + image_version=args.image_version + ) elif "assign_image" == args.command: fun( image_name=args.image_name, image_version=args.image_version, client_mac_address=args.mac_address, ) +elif "detach_image" == args.command: + fun( + image_name=args.image_name, + image_version=args.image_version, + client_mac_address=args.mac_address, + ) elif "run" == args.command: fun() else: diff --git a/network/__pycache__/communication.cpython-310.pyc b/network/__pycache__/communication.cpython-310.pyc index fab1572b98270b1aa95e677e0d147c9b9c88c46c..4f01d07597396bc71f7742e5f411b62ff7051968 100644 GIT binary patch delta 1931 zcmaJ>&u<$=6rS1jt~XwL{o8dygOVQFSX418NUflvA(T>)G;Nc#CCd+Y=pe zl4e=QhxA%5Q*q*uAaUrW5=01zQ*THWa7LDpxNzXYKTwc(Z`X01G+|fs`SaV^_h!ED z&Hgg}j1D_5s+f(9<7FU6i;UIAlap@kQZX)p2Wu-l>viW`AXiV zErqF`QVM&@9WqB8Etn4&9wdO8j{}2bnXt%1rKS2=2RwC`QUaOu1DT?p(%w%MB2Hwp zMoBA7iJOI~F-9ES3vbh2VPn6rN)}*2t?*WOuKT+|H-Ley@fjW%f}2J&-9&k%V%MFK zdczr6YOK|nFMj0_g$XxF8}NnGxA2ANWy`Tf)>fJ2*vxm1e}vh5yP`+hg-&$_mu;sM z6e}%mmTFeHWHH9=X0y*fbij<&8mzLkiIXQ`X?><@N119y8Y`@u*>}^7pvyriaJHC! z?JeANRMMzQC7PfbU~p^ystU%`?hWjDG~EAGBvL2npJFS0(vl&`{-pxxwF>r^8v`%N z4)rAH4%v}Bc}L!b^M-u@=7l!L!Mu=fp2R4ar+6VSPwIA!4CblgNouJ5I5ZOo>K}qj zkQ#c~@AWG1K&z|BdcE$qenEYqO()gWxld$M@5aD4i%(;My9h=PQouU#;Xa#i>Qjg| zhX5ogN)3@rW!r9U&gSbjA3rp;e5^ota$WtfKs|Bf&9jHZf#Y>Z90iY*-`V1aRI<%i zASLV>g)v9_7P4L|tyG%Mku7$-V!VcVba=VKntp=gCCWjy(fv;LYc-Zi>$T>ukfNtl z1eJQ*9e!Dyf_v!EEt2&D?!z0xMi!W2AK7e-Vf$SQ1#0 zDTm@ImWFuhV0V{BI#t5bhyhUEUB^ zvRBpLfi~Y1kFrHGFXd^T8j>673eDH*Gle)G!P%n-rx2z9I&uZRB~X*_a|nJB9of;9 vV!6|&UWLaR8JD!UCQIidRLZ~@f^ix~EJCA(n96;lPhb{*XCH{?x!L~!Wht@h delta 776 zcmZvZ&rcIU6vtTH^GyxM6R1ASutF~%s71yP+Vu=1RI{)mZJEho=j{RcRDForYUI5=<1!5BB0&wS_2?3?#}^J4mY##jr50!Vz} z?q+p=__I+Um~n?VaT=q8&rwCEahecwkeWqgC3}11J;62pQfIi*3k07?rbSp@Bw$wo z*8tZ61^&v&;TitZIHk@*V1b+AGS2o`cpa;kA##QPikyqsAnySjz%oA*J+V~;c>{10 za2hZrVChT&eG5%Hp4nvZ)jg<|c0XSu8cr`W=DG|8cDS?f5Jl@Gbk!e4`5x z1Gkp?KaVU6`%8Q=c2XY^GKe*-Vu?>2HjnPNl0}z2c3Ea)ojRUtS0C2gx@QaR@MnpX z^cC?pi7&yok#`99OEvfha}MYDt~otwMcD;Neh~moVq<_y0RJ`4^X=rAo`dyWzavXSAlsqw5Kw*m(hu#I}&3b7G5mj;Xiz}${pt9$GSEGGeeJ{?p)DX0}5BFECl og&Y>65fji*M$*EXEKLQl6boQ26%6y(&_~_xpI`VK|29Ec*x8SP9h(>7~nJ6p@PbG4k1>#cmdP%DUIqgAXG&r2-B%nu}H-d1WQW?Yu4 z*^e=gRF%DAgWGl2w(QPYvtvV0aBRnEc029*`!>h8B49E z-EqClyKh}+*Vk-Md&jZ4r@cv+WpCz8-K|sCvo{*{ri;Or8#G<~qvPZa{5wCFCCQd* zGLvcwlWmnLwz{e@m1!SnHI3=EZX5fhRfQQW0|^#_K9HeLmgOKZNrLNnR=`=76~%S` zE{&Dg6r?#eZRex0G&aM^kQUe;jH9r<>=B$5=`0v|mhFS26kMNU`*AkK4%pMUp2y4w z*&#?~#E5zEbdR#bkd)bxV0Myftip~$x`!Q0NFQUzA>GR!Pe@O&Cm?-<&0_|IRoO|L z&9W!ObJLtpu~U%jV+-Q?BKsW9=GfC1;~Al`XW4U*?2k!4&z^_m0DHkch!J08FX8ME zds(dFAVz$Jy$Z>r>^1QmBzc{E0g}V)4VcJj&kS+&iY!3XT5Yy$uh{CY!Ccm?mVM1` zHOP3wkAVX+^(Fi}zd_S~(S}?29l)=qzz?59RP2peyGTjt3j$m zs&5z@8OXH{Vg9Z;%(}TdGOhtXvH@li?kxXH<0kId$PWtxZJ_Lz?x6P-xXWcpliXtO zvMk+%oAjsL(m=mWBTcaaeP5B@L$90qEuH%MN7v=+Qnk3;&pB6HP1k;TadGinRq?cy zW`}vx_KnRZx2&JVkIZVwwHP4`_8lJIy@$JiJ-}dw~ zXWxA1yr*4QK6lZ}TseJd`NHyfPd|I<(#1=jxpe8m+ZUEjf92zCe5%u2!OfgB>SZqb zo>77{3`Em*bUcNTun<&AaQ zaS>&@{1gUV+;W>OXEC(8=wM;9N2_UenyzJ$>K#(JB9|0Xo|X;#7gR-2HcDHz^L12a-4F{18{|F~*(oMO4d>{>_JJL;MpbX_Z(obe& z2}g>n-jRnIGhQHeg)mNl}d|7&0a*kYw&C%a``1hr-PrY{r_KAM|GHlowX5b+z zc?elC%@$axdP}8oV5K+`R;n!bPo3qw%PYXwg$gUJrsGz+tCfa`$xdaJciWXvb1zq_ zCO?QJc>0PBm-Y0kTQ=`;3csETgW-?ih#wU_Y#O(!@e_0g#W-6n#P}c|aMHAZ*KbZz zr#v3Q3MM`VfkRH9LH(wjBQUl#z5J=Bm+qJ4lB^L3_xpR#1hHfxcC(jn*Be%yG0f^T z$Opm?J!u5SFK_~7Ab1EUR8pHI_fHJep#n(u0TOknx*8zSU2Ve{$nQg?!wkVT29F}( znID4(Qov*Oz6^LYz%vbavcnux23iPD0F8DFVMjoNBZB6xx_qwxcoG|5HI{0S++u{0cD>F!Rb4? z*jnkqVIw8ew9auffCZkT05DT203_``U^C1A)XdU@Lq2E^`o~g08yhfzS_T70sQnfv zDX2Xcp!WDKsA+Hp=zhaIBYn%fsl(y21APMtL6)u_O-1ziR7CT=_kgGoAX*p}fv7Qt zsBy~(5XBK8%0G_>ua+ia#h=G@AB7iis~FsQ3e7ma%6(1G!5x+D!oRe9`5#(7hV#iZ zoYN=pO8zri*KAINF@b-=Yl9G~z)u6e1@aU^_Ym?M2%!U+$sZ~rjHdhP(Nt9ZFdF(P zh0%xb)2h0R4E_c*(_c*WRuNF)vQBr^g}d36V1ZFA9dlU@LNATO8n`G;^B-LF(Gipm zxl~fC%NW2pgeHutdnu}F<qt>R^>yAas3NLv%-_;;r=dYm4I5dXuac|Nkik`KC zuX0B?FMo~XX~}Ni7*L@uzm6MXiO-Y^h^PwqhGrxZUu6Q1TqOG@tT90>-TOmVh2HlHyKA;+hQH{rjlVA|9eH_mBFR zA*&G$K|R=Xdlk1E;~xDZ|Dns>tRGw{8G(*4hK$$Ra@G<51l5i~_c?+ki{`FRP<)&w zCcYd73tWleCT9+fPfL*Jkn2v|FKJ~Q30}A-LS}SV0*6F6`#py0pB=;D86O9q4Z7-v z4o}eq`7rK4YdAwqV`cDd;@iY`7T;MQl^cLTio5fK86*YqtzMCo1fvw_y`G9*vG;AD zM}d&gD-EaMdl6xArz7+Ns)7*MpZB3@cUg0_M*wzLh^E8ED(<=+2l3c=&j4!lWC4|{ z&WvyzahdV_m}_o^6wfo)yN*j$iHO9Wjsp}X=b>8uMQS3-!z18B+9Mp4fcJ9U7PI{4 zPW9NLzq-HEY?qV@In;K4q`t&0oJdJRhjGWc_fuT$KRKquxt)#`bE|YG zIi9N3Oz(W)B&1brEX18;#OedBhOUBwD{vKrf1y;7UWDoIP;73bSlP=e&KW(f>d4SB9%EmLm!9sd3-g4KFjkDN@}3wQ34s=5%LfU zauaQyU!u9+K~vqucCJvTv{eKwQ(Pz8;S-lK6FnGlOnRn|(Er{6H2#YG;!nA+=<}bH zGQqkIq#F<2?(O%DkdDI=z(TExRPx`g)|*Qi7<6EzyPNyeX!X0RfObqqo|E$2zrh%S z$SEpO1QQ!cNDSMWs(h+sC_MyPsZQ3)m?roVhTL5pRT%j*yBfJ5zs{!JXs$MGRv}&s z`n!v{(4@QRN?hMvS4nfBDStMcrwi;7D^8@zAJYK;I}Uqpa*!bo>+iAc zWU@~zb+=lq^MqS*>euXGC(V!WEWu|hQAZHPidbus@$$1+U>ID!O)3+5P7CB&;+kTJ zKcL^ebb^Y_OWL?BM-~^|w!7@6eywfcw%*in4teG1$6LER0?7P4I5_QrPwY* zA%1{@JEZCmn_7a6szcppp+oRjO2MhY{;AYppXptkoM_!#n05mkVu|*6q+XwO#W*Po zn5Nh~dz%z15+-eJ`r-a9I(jJnSf`#;S5L;K5U@T0&!5oM4-2BQ%VAH9Y6!s@KP?%w zwJAV)BmyaZw>KpQvD~iZbD;!G3Fr7VCi)*b$ECnI4n+kXijj}-e$hoKa+_YzXav0g zXvdYEPV=ZAw8WyQe(=3P1gj={t*;2AM)-V{6h09j0@G8GK17KZI#7e|ndm@;fLn@* z|AMYg1&yE0Mn#?>@0B%BB5DECXZ;aO@%z1Uuq^!1wY40R{R8wcxNwUYJ@! zET#HjskZ^RX~T0b;GP;zGufRP7y@SD;chuD3E&MjJ}D~R2w@GP0u<3TEwIkk@dA6T6U0XNo&U^# z!x+v2nlQW5X18+RG)7szkO}PO@1>bp8O0#Q-XoD)r;j_)K=>LaBUR2vs^oUNiv;7l zig;1;U+a@z)X<<0t5UvBx}?{1vG+K04DS*f>w(JX!2+e3C=wjf`NJwrL5viR_Mm?> zQrx7`J(F_S7Al>|LGk~}7=pKQW5I8SchCxm2P8L65?r7~3ZEFtOiJrC%5mtMe+PYI z?<&JICt*5Qse4+;n~pypJ=aTQM-IC8_ouFb`YczE@Cv$kIm-fL*=o@* z1PYdQb*tVI)se+TAbyhiXcY0q8%nY5&R>F%e}S6Q)QA{=hK{~O%|&XyhQ?FzlL9XI z+Pfs9y>k9dYQ9B{Ma>E|HZ@IZ-ls-j-=-r@%@#GIm?e%*l8lmZ7XQvUG^SLR%X8&? zX==7Sn=2J?G+j2!`Pr;e#4iH24JG$rcBWK>yrfK#1aCdJl>^~X=@9zn#F*8o8vPig z-ooz^+CA~3j9O_VSm^S$_>o5KKul=ga2H~~mf}S2yglNl3=2@O-Ex92Uynj(VxcBF y){gld?eI5K&O(S>^b-wX@BA=&i-%0ph-bd)V~HI^sHKFaAQcdl=LC5ztN#XY<;Ir) delta 3236 zcma)8TTC2P7@jjbJG;#84lpcxy=> z;$qrF6JtzzUQC-Npbt$=qfw$y+F0L=4}H-3urbjXjV3;sJ~U}d@&C_WfMr2u5btE9ka$)CSrJR5xt`6(Vh>VOUkt2sa#+Zl^6k z@z5P0wNtkMrIqdiikG(OK9Fdq9k6Qg)7{)oA24*%E+7Ty9fza>UUfTiSWcRfVb2mz$6x?xNi5y@y|pURcTn{gwYhX}GbM)kzt z@I*|HCrlLNAQ6$+i9*VVb48N1dy=d^kbxqtsG>yd^WuY9H;X#mtj^Y4dl96lFefI3 z1z}N`v&=RnEem8(c;8Eeq=h|E;brG-9)IIGLInJtfnUFn6#E28NZ4kMvVQ^rmT>q* zv6ua23+svzKofT$>_o5vq~(#=L;~E|jSK5Y%rFzis-x1ms1c7?jYi~07NZTJ9if=l ze)hfHzpWLv%?1GQ2|QK^`Yo=MfuG>FRl6@)*_ZZK_M>uxbTj7I+grpk15a5s#n=PO zoF(13wYeqC=d5qv2RtSwB^PuF%#c>_mk<6j247u~LFHe zEmSXO0rm46*kp86AB@q99{=Gx?||5BbGE<4q{zF=4c#SC9Va4INgTG7)kKt3X383U{Gv|Qkjt?MXmI^~4{7|)E&uJ^% zqKY2&TV?lwXMrceW7mhro?LTVlG1&^VoOZ zHt7-IgzAi#7vN#Gv7kD&_0@@$$$oOTluI)Qv+Z04?QC6;**7-NYE9yytk%~xVJljb z#mf&Uv`|-Bz#|$HTcKt zn*9o+>7yFA^dbnbpM!p4F5RxNlGLh^`774wME!h-@*?6zg>q@_jY)oS3{J%dC!@n- zvAAwLip*O&n6|(b&BuSlmg&e>Q3g*mf;_M|bqyZp!VYDzU{&xYY@4+JnN|sq062?Y zDx~S2YGX^m2=mtcvTtqHhVaNuvu4Mvy$;jr7P7s>cGcf4jD!t$cqAx36=|o8OWBs^ zOsFp{;kiXejf%yGyI6ZMng;GWdAG{BM#|YmA+12wAF)AiLC1n|cjN+mT(umk9n=W^^EtU%vPn zWAKuB9G`j+o}o69#XQlSq3Wp&#Ki!T z1pr?!`0JBh(1_;KoF0eAsk$_Ws%myk@hGr%sEX#$N`V5;?0Vy3pe852<~HW{a)gaG GHT(;XWOttc diff --git a/utils/database/database.py b/utils/database/database.py index 42de3ce..da85caa 100644 --- a/utils/database/database.py +++ b/utils/database/database.py @@ -24,7 +24,7 @@ class Database: "INFO": 20, "WARNING": 30, "ERROR": 40, - "CRITICAL": 50 + "CRITICAL": 50, } self.logger.setLevel(log_level_mapping_dict[logging_level]) except Exception as ex: @@ -38,18 +38,39 @@ class Database: with session.begin(): result = session.query(Client).all() except Exception as ex: - self.logger.error( - f"Error getting list of clients from database: {ex}") + self.logger.error(f"Error getting list of clients from database: {ex}") result = [] return result def get_client_by_mac_address(self, mac_address: str) -> Client: result = None session = self.Session() + session.expire_on_commit = False + try: + with session.begin(): + result = ( + session.query(Client) + .filter(Client.mac_address == mac_address) + .first() + ) + except Exception as ex: + self.logger.warn(f"Error getting client by mac address: {ex}") + return result + + def get_client_vm_list_by_mac_address(self, mac_address: str): + result = None + session = self.Session() + session.expire_on_commit = False try: with session.begin(): - result = session.query( - Client).filter(Client.mac_address==mac_address).first() + client = ( + session.query(Client) + .filter(Client.mac_address == mac_address) + .first() + ) + result = [] + for vm in client.vm_list_on_machine: + result.append(vm.image_id) except Exception as ex: self.logger.warn(f"Error getting client by mac address: {ex}") return result @@ -59,11 +80,9 @@ class Database: try: session = self.Session() with session.begin(): - result = session.query( - Client, client_version=client_version).all() + result = session.query(Client, client_version=client_version).all() except Exception as ex: - self.logger.warn( - f"Error getting client list by software version: {ex}") + self.logger.warn(f"Error getting client list by software version: {ex}") return result def get_clients_by_vm_image(self, vm_image: VMImage) -> list[Client]: @@ -74,8 +93,7 @@ class Database: if client.has_vm_installed(vm_image.image_hash): result.append() except Exception as ex: - self.logger.warn( - f"Error getting list of clients with VM installed: {ex}") + self.logger.warn(f"Error getting list of clients with VM installed: {ex}") result = [] return result @@ -92,10 +110,12 @@ class Database: def modify_client(self, client: Client) -> Client: try: - old_object = self.get_client_by_mac_address(client.mac_address) session = self.Session() with session.begin(): - old_object = client + old_object: Client = session.query(Client).filter(Client.mac_address==client.mac_address).first() + old_object.ip_address = client.ip_address + old_object.hostname = client.hostname + old_object.client_version = client.client_version session.merge(old_object) session.flush() session.commit() @@ -116,8 +136,9 @@ class Database: try: session = self.Session() with session.begin(): - response = session.query( - VMImage, image_id=image_id).first() + response = ( + session.query(VMImage).filter(VMImage.image_id == image_id).first() + ) return response except Exception as ex: self.logger.error(f"Error getting image data from database: {ex}") @@ -129,40 +150,50 @@ class Database: response = session.query(VMImage).all() return response except Exception as ex: - self.logger.error( - f"Error getting list of images from database: {ex}") + self.logger.error(f"Error getting list of images from database: {ex}") def get_image_by_name(self, image_name: str) -> list[VMImage]: try: session = self.Session() with session.begin(): - response = session.query( - VMImage, image_name=image_name).all() + response = ( + session.query(VMImage) + .filter(VMImage.image_name == image_name) + .all() + ) return response except Exception as ex: - self.logger.error( - f"Error getting list of images from database: {ex}") - - def get_image_by_name_version_string(self, image_name_version_string: str) -> list[VMImage]: + self.logger.error(f"Error getting list of images from database: {ex}") + + def get_image_by_name_version_string( + self, image_name_version_string: str + ) -> list[VMImage]: try: session = self.Session() with session.begin(): - response = session.query( - VMImage).filter(VMImage.image_name_version_combo==image_name_version_string).first() + response = ( + session.query(VMImage) + .filter( + VMImage.image_name_version_combo == image_name_version_string + ) + .first() + ) return response except Exception as ex: - self.logger.error( - f"Error getting list of images from database: {ex}") + self.logger.error(f"Error getting list of images from database: {ex}") def get_image_by_hash(self, image_hash: str) -> list[VMImage]: try: session = self.Session() with session.begin(): - response = session.query(VMImage, image_hash=image_hash) + response = ( + session.query(VMImage) + .filter(VMImage.image_hash == image_hash) + .first() + ) return response except Exception as ex: - self.logger.error( - f"Error getting list of images with specified hash: {ex}") + self.logger.error(f"Error getting list of images with specified hash: {ex}") def add_image(self, image: VMImage): try: @@ -187,15 +218,41 @@ class Database: return old_object except Exception as ex: self.logger.error(f"Couldn't modify object in database: {ex}") + raise DatabaseException(f"Couldn't modify object in database: {ex}") + + def delete_image(self, image_to_delete: VMImage): + try: + session = self.Session() + with session.begin(): + session.delete(image_to_delete) + session.flush() + session.commit() + except Exception as ex: + self.logger.error( + f"Error deleting image with id={image_to_delete.image_id}: {str(ex)}" + ) raise DatabaseException( - f"Couldn't modify object in database: {ex}") + f"Error deleting image with id={image_to_delete.image_id}: {str(ex)}" + ) - def assign_image_to_client(self, client_mac_address: str, image_name_version_combo: str): + def assign_image_to_client( + self, client_mac_address: str, image_name_version_combo: str + ): try: session = self.Session() with session.begin(): - client = session.query(Client).filter(Client.mac_address==client_mac_address).first() - image = session.query(VMImage).filter(VMImage.image_name_version_combo==image_name_version_combo).first() + client = ( + session.query(Client) + .filter(Client.mac_address == client_mac_address) + .first() + ) + image = ( + session.query(VMImage) + .filter( + VMImage.image_name_version_combo == image_name_version_combo + ) + .first() + ) client.vm_list_on_machine.append(image) session.merge(client) session.flush() @@ -204,6 +261,34 @@ class Database: self.logger.error(f"Couldn't add image to client list: {str(ex)}") raise DatabaseException(f"Couldn't add image to client list: {str(ex)}") + def detach_image_from_client( + self, client_mac_address: str, image_name_version_combo: str + ): + try: + session = self.Session() + with session.begin(): + client = ( + session.query(Client) + .filter(Client.mac_address == client_mac_address) + .first() + ) + image = ( + session.query(VMImage) + .filter( + VMImage.image_name_version_combo == image_name_version_combo + ) + .first() + ) + client.vm_list_on_machine.remove(image) + session.merge(client) + session.flush() + session.commit() + except Exception as ex: + self.logger.error(f"Couldn't remove image from client list: {str(ex)}") + raise DatabaseException( + f"Couldn't remove image from client list: {str(ex)}" + ) + def add_user(self, new_user: User): try: session = self.Session() @@ -219,8 +304,7 @@ class Database: try: session = self.Session() with session.begin(): - user = session.query(User).filter( - User.user_id == user_id).first() + user = session.query(User).filter(User.user_id == user_id).first() return user except Exception as ex: self.logger.error(f"Error getting data from database: {ex}") @@ -230,8 +314,7 @@ class Database: try: session = self.Session() with session.begin(): - user = session.query(User).filter( - User.username == username).first() + user = session.query(User).filter(User.username == username).first() return user except Exception as ex: self.logger.error(f"Error getting data from database: {ex}")