#!/usr/bin/env python from snack import * import gettext, sys, os, re, string, time import partition_data as p_d screen = SnackScreen() screen.drawRootText(0, 0, "Welcome to Rocky Security OS") screen.pushHelpLine( (" F1-New F2-Delete F3-Boot F4-Raid F12-OK ")) global count count = 0 def DEBUG(): global count print 'debug' + str(count) count = count + 1 global base_unit base_unit = 'B' # the origin unit is 's' #def pretty_unit(num): # if num*0.5*1.024 < 10*1000: # # < 10M unit is kB # num_pretty = num*0.5*1.024 # unit = 'kB' # elif num*0.5*1.024 < 10*1000*1000: # # < 10G unit is MB # num_pretty = num*0.5*1.024/1000 # unit = 'MB' # elif num*0.5*1.024 < 10000000000: # # < 10T unit is GB # num_pretty = num*0.5*1.024/1000/1000 # unit = 'GB' # return str("%.0f" %num_pretty) + unit # #def unit_to_s(num): # if re.search('kB', num): # # unit kB to s # num_s = num/1.024*2 # elif re.search('MB', num): # # unit MB to s # num_s = num/1000/1.024*2 # elif re.search('GB', num): # # unit GB to s # num_s = num/1000/1000/1.024*2 # return num_s #the origin unit is 'byte' def pretty_unit(num): num = int(num) global num_pretty global unit_pretty if num < 1000: # < 1k unit is B num_pretty = num unit_pretty = 'B' elif num/1000 < 10*1000: # < 10M unit is kB num_pretty = num/1000 unit_pretty = 'kB' elif num/1000/1000 < 10*1000: # < 10G unit is MB num_pretty = num/1000/1000 unit_pretty = 'MB' elif num/1000/1000/1000 < 10*1000: # < 10T unit is GB num_pretty = num/1000/1000/1000 unit_pretty = 'GB' return str("%.1f" %num_pretty) + unit_pretty def unit_to_byte(num_unit): #num = int(filter(str.isdigit, str(num_unit))) # because int('10.0') error num = int(float(get_size_num(str(num_unit)))) if re.search('kB', num_unit): # unit kB to byte num_byte = num*1000 elif re.search('MB', num_unit): # unit MB to byte num_byte = num*1000*1000 elif re.search('GB', num_unit): # unit GB to byte num_byte = num*1000*1000*1000 return num_byte def sort_partition(dev): list_partition = [] list_partition_sorted = [] for partition in p_d.Partition.dict[dev]['partition'].keys(): list_partition.append(partition) list_partition_sorted.append(partition) print list_partition for partition in list_partition: sn = p_d.Partition.dict[dev]['partition'][partition]['sn'] list_partition_sorted[int(sn)-1] = partition return list_partition_sorted def showPartitionWindow(): global clb global bb bb = ButtonBar (screen, (("New", "new", "F1"),("Delete", "delete", "F2"),("Boot", "boot", "F3"),\ ("Raid", "raid", "F4"),("OK", "ok", "F12"),("Cancel", "cancel"),)) clb = CListbox(height=10, cols=6, col_widths=[19,11,11,11,10,12], scroll=1, returnExit = 1, width=85, col_pad=2, col_labels=['Device', 'Start', 'End', 'Size', 'Type', 'Flags'], col_label_align=[CENTER,CENTER,CENTER,CENTER,CENTER,CENTER]) for dev in sorted(p_d.Partition.dict.keys()): # dev: /dev/sda dev_name: sda; dev: /dev/cciss/c0d0 dev_name:cciss/c0d0 dev_name = dev.split('/dev/')[1] partition_table = p_d.Partition.dict[dev]['partition_table'] dev_size = p_d.Partition.dict[dev]['disk_size'] #dev_size_pretty = pretty_unit(filter(str.isdigit, dev_size)) clb.append(["%s" %(dev),"","","","",""], \ "dev:%s:%s" %(dev,dev_size), [LEFT, RIGHT, RIGHT, RIGHT, LEFT, LEFT]) if partition_table == None: continue list_partition_all_sorted = p_d.sort_partitions(p_d.Partition.dict, dev, "all") for partition in list_partition_all_sorted: start = p_d.Partition.dict[dev]['partition'][partition]['start'] end = p_d.Partition.dict[dev]['partition'][partition]['end'] size = p_d.Partition.dict[dev]['partition'][partition]['size'] ptype = p_d.Partition.dict[dev]['partition'][partition]['type'] flags = p_d.Partition.dict[dev]['partition'][partition]['flags'] try: if int(filter(str.isdigit, start)) < 2000000 and \ int(filter(str.isdigit, size)) < 2000000: continue except: pass start_pretty = pretty_unit(filter(str.isdigit, str(start))) end_pretty = pretty_unit(filter(str.isdigit, str(end))) size_pretty = pretty_unit(filter(str.isdigit, str(size))) # pretty free space partition if re.search('free', partition): partition_pretty = 'Free Space' else: partition_pretty = partition clb.append([" %s" %(partition_pretty),"%s" %(start_pretty),"%s" %(end_pretty),"%s" %(size_pretty),"%s" %(ptype),"%s" %(flags)],\ "partition:%s:%s:%s:%s:%s:%s:%s" %(partition, dev, size, start, end, partition_table, size_pretty),\ [LEFT, RIGHT, RIGHT, RIGHT, LEFT, LEFT]) sg = Grid(2, 3) sg.setField(bb, 0, 1, anchorLeft = 1) g = Grid(1, 3) g.setField(clb, 0, 0) g.setField(sg, 0, 2) g.place(1, 1) screen.gridWrappedWindow(g, "Partitioning") f = Form() f.add(bb) f.add(clb) f.addHotKey("F1") f.addHotKey("F2") f.addHotKey("F3") f.addHotKey("F4") f.addHotKey("F12") global res res =f.run() screen.refresh def get_size_num(size): if re.search('kB', size): size_num = size.split('kB')[0] elif re.search('MB', size): size_num = size.split('MB')[0] elif re.search('GB', size): size_num = size.split('GB')[0] return size_num def show_error(tips): g21 = GridForm(screen, "Error", 1, 2) g21.add(Label("%s." %(tips)), 0, 0) g21.add(Button("OK"), 0, 1) g21.runOnce() def check_size_set(size, maxsize): if size == "": return ("Must specify a value") try: if float(size) > float(maxsize): return ("Requested value is too large") except: return ("Requested value is not an number") def count_primary_extended(dev): count_primary = 0 count_extended = 0 for p in p_d.sort_partitions(p_d.Partition.dict, dev): if p_d.Partition.dict[dev]['partition'][p]['type'] == 'primary': count_primary += 1 if p_d.Partition.dict[dev]['partition'][p]['type'] == 'extended': count_extended += 1 return (count_primary, count_extended) def get_extended_start_end(dev): extended_start = '' extended_end = '' for p in p_d.sort_partitions(p_d.Partition.dict, dev): if p_d.Partition.dict[dev]['partition'][p]['type'] == 'extended': extended_start = p_d.Partition.dict[dev]['partition'][p]['start'] extended_end = p_d.Partition.dict[dev]['partition'][p]['end'] return (extended_start, extended_end) def add_partition(ptype='', extended='', label=''): free_space = clb.current().split(":")[1] dev = clb.current().split(":")[2] dev_name = dev.split("/dev/")[1] size = clb.current().split(":")[3] size_pretty = clb.current().split(":")[7] free_space_start = filter(str.isdigit, clb.current().split(":")[4]) free_space_end = filter(str.isdigit, clb.current().split(":")[5]) free_space_sn = int(p_d.Partition.dict[dev]['partition'][free_space]['sn']) # 4k assign #if int(free_space_start) < 2*1000*1000: # free_space_start_real = free_space_start # free_space_start = '0%' size_num = filter(str.isdigit, size) #size_num = get_size_num(size) unit = filter(str.isalpha, size_pretty) #size_pretty_num = filter(str.isdigit, size_pretty) size_pretty_num = get_size_num(size_pretty) bb2 = ButtonBar(screen, (("OK", "ok"), ("Cancel", "cancel"))) sg21 = Grid(2,1) typeLbl = Label("File System type: ") lb_fstype = Listbox(height=3, scroll=1) list_fs = ['ext2', 'ext3', 'ext4', 'linux-swap', 'raid', 'reiserfs', 'xfs', 'jfs'] for fs in list_fs: lb_fstype.append(fs, fs) lb_fstype.setCurrent('ext3') sg21.setField(typeLbl, 0, 0, anchorLeft=0) sg21.setField(lb_fstype, 1, 0) sg22 = Grid(2,1) sizeLbl = Label("Size(%s): " %(unit)) entry_size = Entry(10, size_pretty_num) sg22.setField(sizeLbl, 0, 0) sg22.setField(entry_size, 1, 0) sg23 = Grid(2,1) rb_primary = SingleRadioButton("primary", None, 1) rb_extended = SingleRadioButton("extended", rb_primary) sg23.setField(rb_primary, 0, 0) sg23.setField(rb_extended, 1, 0) cb_boot = Checkbox("Force to be a bootable partition") g2 = GridForm(screen, "New Partition(%s)" %(dev), 1, 9) g2.add(sg21,0,0) g2.add(Label(""),0,1) g2.add(sg22,0,2) g2.add(Label(""),0,3) g2.add(cb_boot,0,4) g2.add(Label(""),0,5) if extended == "extended" or label == 'gpt': g2.add(bb2, 0, 6) res2 = g2.runOnce() if label == 'gpt': ptype = 'primary' else: g2.add(sg23,0,6) g2.add(Label(""),0,7) g2.add(bb2, 0, 8) res2 = g2.runOnce() if rb_primary.selected(): ptype = 'primary' else: ptype = 'extended' if bb2.buttonPressed(res2) == "ok": #if ptype == 'extended' and count_extended != 0: # show_error('An extended partition has existed !You can not create any more !') fs_set = lb_fstype.current() tips = check_size_set(entry_size.value(),size_pretty_num) if tips: show_error(tips) else: size_set_pretty = entry_size.value() + unit size_set_num = unit_to_byte(size_set_pretty) size_set = str(size_set_num) + base_unit partition_start = free_space_start + base_unit partition_end = str(int(free_space_start) + size_set_num) + base_unit flags = '' if cb_boot.value() == 1: if flags == '': flags = 'boot' else: flags = flags + ',boot' if fs_set == 'raid': if flags == '': flags = 'raid' else: flags = flags + ',raid' partition_num = 0 list_partition_sorted = p_d.sort_partitions(p_d.Partition.dict,dev,'partition') list_free_sorted = p_d.sort_partitions(p_d.Partition.dict,dev,'free') if list_partition_sorted == []: # an empty disk pass else: if ptype == 'logical': partition_num = 4 for p in list_partition_sorted: if int(p_d.Partition.dict[dev]['partition'][p]['num']) > partition_num \ and p_d.Partition.dict[dev]['partition'][p]['type'] =='logical': partition_num = int(p_d.Partition.dict[dev]['partition'][p]['num']) else: if label == 'msdos': for p in list_partition_sorted: if int(p_d.Partition.dict[dev]['partition'][p]['num']) > partition_num \ and p_d.Partition.dict[dev]['partition'][p]['type'] != 'logical': partition_num = int(p_d.Partition.dict[dev]['partition'][p]['num']) elif label == 'gpt': if int(p_d.Partition.dict[dev]['partition'][list_partition_sorted[0]]['num']) != 1: pass else: for i in range(len(list_partition_sorted)): #print int(p_d.Partition.dict[dev]['partition'][sorted(list_partition)[i]]['num']) if i == len(list_partition_sorted) - 1 or \ int(p_d.Partition.dict[dev]['partition'][list_partition_sorted[i]]['num']) \ != int(p_d.Partition.dict[dev]['partition'][list_partition_sorted[i+1]]['num']) - 1: partition_num = int(p_d.Partition.dict[dev]['partition'][list_partition_sorted[i]]['num']) break partition_num += 1 partition = dev + str(partition_num) partition_sn = free_space_sn if size_set_pretty == size_pretty: partition_end = free_space_end if len(list_free_sorted) > 1: for p in list_free_sorted: if int(p_d.Partition.dict[dev]['partition'][p]['sn']) >= int(free_space_sn): p_num = int(p_d.Partition.dict[dev]['partition'][p]['num']) p_d.Partition.dict[dev]['partition'][p]['sn'] = p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['sn'] p_d.Partition.dict[dev]['partition'][p]['start'] = p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['start'] p_d.Partition.dict[dev]['partition'][p]['end'] = p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['end'] p_d.Partition.dict[dev]['partition'][p]['size'] = p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['size'] p_d.delete_one_partition(p_d.Partition.dict, dev, 'free '+str(len(list_free_sorted)+1)) p_d.delete_one_partition(p_d.Partition.dict, dev, 'free '+str(len(list_free_sorted))) else: del p_d.Partition.dict[dev]['partition'][free_space] else: list_partition_all_sorted = p_d.sort_partitions(p_d.Partition.dict, dev, 'all') p_d.Partition.dict[dev]['partition'][free_space]['start'] = str(int(free_space_start) + size_set_num + 1) + base_unit for p in list_partition_all_sorted: if p_d.Partition.dict[dev]['partition'][p]['sn'] > free_space_sn-1: p_d.Partition.dict[dev]['partition'][p]['sn'] = int(p_d.Partition.dict[dev]['partition'][p]['sn']) + 1 p_d.Partition.dict[dev]['partition'][free_space]['sn'] = int(free_space_sn) + 1 p_d.Partition.dict[dev]['partition'][free_space]['size'] = \ str(int(filter(str.isdigit,p_d.Partition.dict[dev]['partition'][free_space]['end'])) - \ int(filter(str.isdigit,p_d.Partition.dict[dev]['partition'][free_space]['start']))) + base_unit p_d.Partition.dict[dev]['partition'][partition]['sn'] = int(partition_sn) p_d.Partition.dict[dev]['partition'][partition]['num'] = int(partition_num) p_d.Partition.dict[dev]['partition'][partition]['start'] = partition_start p_d.Partition.dict[dev]['partition'][partition]['end'] = partition_end p_d.Partition.dict[dev]['partition'][partition]['size'] = size_set p_d.Partition.dict[dev]['partition'][partition]['type'] = ptype p_d.Partition.dict[dev]['partition'][partition]['filesystem'] = fs_set p_d.Partition.dict[dev]['partition'][partition]['flags'] = flags if ptype == 'extended': list_partition_all_sorted = p_d.sort_partitions(p_d.Partition.dict,dev, 'all') list_free_sorted = p_d.sort_partitions(p_d.Partition.dict,dev,'free') extended_free_space_sn = partition_sn + 1 free_num = 0 if len(list_partition_all_sorted) == partition_sn: if list_free_sorted != []: for f in list_free_sorted: if p_d.Partition.dict[dev]['partition'][f]['num'] > free_num: free_num = p_d.Partition.dict[dev]['partition'][f]['num'] else: for p in list_partition_all_sorted[::-1]: p_num = p_d.Partition.dict[dev]['partition'][p]['num'] if p_d.Partition.dict[dev]['partition'][p]['sn'] > partition_sn: if re.search('free', p): p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['sn'] = p_d.Partition.dict[dev]['partition'][p]['sn']+1 p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['num'] = p_num + 1 p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['start'] =\ p_d.Partition.dict[dev]['partition'][p]['start'] p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['end'] =\ p_d.Partition.dict[dev]['partition'][p]['end'] p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['size'] =\ p_d.Partition.dict[dev]['partition'][p]['size'] p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['type'] = ' ' p_d.Partition.dict[dev]['partition']['free '+str(p_num+1)]['flags'] = ' ' else: p_d.Partition.dict[dev]['partition'][p]['sn'] = p_d.Partition.dict[dev]['partition'][p]['sn']+1 if p_d.Partition.dict[dev]['partition'][p]['sn'] == partition_sn+1: p_d.delete_one_partition(p_d.Partition.dict, dev, p) if list_free_sorted != []: for f in list_free_sorted: if p_d.Partition.dict[dev]['partition'][f]['sn'] < extended_free_space_sn: if p_d.Partition.dict[dev]['partition'][f]['num'] > free_num: free_num = p_d.Partition.dict[dev]['partition'][f]['num'] free_num += 1 p_d.Partition.dict[dev]['partition']['free '+str(free_num)]['sn'] = int(partition_sn+1) p_d.Partition.dict[dev]['partition']['free '+str(free_num)]['num'] = int(free_num) p_d.Partition.dict[dev]['partition']['free '+str(free_num)]['start'] = partition_start p_d.Partition.dict[dev]['partition']['free '+str(free_num)]['end'] = partition_end p_d.Partition.dict[dev]['partition']['free '+str(free_num)]['size'] = size_set p_d.Partition.dict[dev]['partition']['free '+str(free_num)]['type'] = '' p_d.Partition.dict[dev]['partition']['free '+str(free_num)]['flags'] = '' def del_partition(partition,dev): bb3 = ButtonBar(screen, (("Cancel", "cancel"), ("Delete", "delete"))) g3 = GridForm(screen, "Del Partition", 1, 3) g3.add(TextboxReflowed(45, 'You are about to delete the partition "%s".' \ %(partition)), 0, 0) g3.add(Label(""),0,1) g3.add(bb3, 0, 2) res3 = g3.runOnce() if bb3.buttonPressed(res3) == "delete": list_partition_all_sorted = p_d.sort_partition(p_d.Partition.dict, dev, "all") partition_sn = int(p_d.Partition.dict[dev]['partition'][partition]['sn']) for p in list_partition_all_sorted: p_sn = p_d.Partition.dict[dev]['partition'][p]['sn'] if p_sn == partition_sn -1 and re.search('free', p): p_d.Partition.dict[dev]['partition'][p]['end'] p_d.Partition.init_from_os() while(1): global ret showPartitionWindow() if bb.buttonPressed(res) == "cancel": ret = 1 break if bb.buttonPressed(res) == "ok": ret = 0 break # assign label(msdos or gpt) of dev if re.search("dev:", clb.current()) and bb.buttonPressed(res) == "new": dev = clb.current().split(":")[1] dev_size = clb.current().split(":")[2] rb = RadioBar(screen, (("gpt", "gpt", 0),("msdos", "msdos", 1))) bb1 = ButtonBar(screen, (("Cancel", "cancel"), ("OK", "ok"))) g = GridForm(screen, "Partition table", 1, 4) g.add(TextboxReflowed(25, 'Warning! A new partition table will be created on "%s" and the data on it will be cleared.' \ %(dev)), 0, 0) g.add(Label(""),0,1) g.add(rb,0,2,(0,0,0,1)) g.add(bb1, 0, 3) res1 = g.runOnce() if bb1.buttonPressed(res1) == "ok": label = rb.getSelection() p_d.Partition.dict[dev]['partition_table'] = label p_d.delete_all_partitions(p_d.Partition.dict, dev) # add partition when dev's label is msdos elif re.search("partition:free", clb.current()) and \ clb.current().split(':')[6] == 'msdos'and \ bb.buttonPressed(res) == "new": dev = clb.current().split(":")[2] dev_name = dev.split("/dev/")[1] free_space_start = clb.current().split(":")[4] free_space_end = clb.current().split(":")[5] free_space_size = clb.current().split(":")[3] (count_primary, count_extended) = count_primary_extended(dev) # add logical partition when an extended partition exists if re.search('extended', str(p_d.Partition.dict[dev])): (extended_start, extended_end) = get_extended_start_end(dev) if int(filter(str.isdigit,free_space_start)) >= int(filter(str.isdigit,extended_start)) \ and int(filter(str.isdigit,free_space_end)) <= int(filter(str.isdigit,extended_end)): ptype = 'logical' #print 'ptype' #time.sleep(5) add_partition(ptype, 'extended', 'msdos') elif int(filter(str.isdigit,free_space_start)) >= int(filter(str.isdigit,extended_end)): show_error("Bad partition scheme ! This free space can't be used any more !") else: if count_primary == 3: show_error('You can not add new primary or extended partition any more !') else: ptype = 'primary' add_partition(ptype, 'extended', 'msdos') else: if count_primary == 4: show_error("You can't add partitions any more !") else: add_partition(ptype='', extended='', label='msdos') elif re.search("partition:free", clb.current()) and \ clb.current().split(':')[6] == 'gpt'and \ bb.buttonPressed(res) == "new": dev = clb.current().split(":")[2] free_space_start = clb.current().split(":")[4] free_space_end = clb.current().split(":")[5] free_space_size = clb.current().split(":")[3] add_partition(ptype='primary', extended='', label='gpt') elif re.search("partition:free", clb.current()) == None and \ re.search("dev:", clb.current()) == None and \ bb.buttonPressed(res) == "delete": dev = clb.current().split(":")[2] partition = clb.current().split(":")[1] del_partition(partition, dev) screen.finish() #p_d.pretty(p_d.Partition.dict) p_d.write_to_xml_file()