set executable bit
[mpi-image.git] / scripts / init.py
1 #!/usr/bin/python
2
3 import sys
4 import os
5 import signal
6 import socket
7 import fcntl
8 import pdb
9 import getpass
10
11 INSANE = 500
12 BLOCKING_DELAY = 30
13
14
15 def alarm(signum, junk):
16     # Do nothing, default implementation does needful
17     return
18
19 def fetch_output(infd,outfile):
20     try:
21         os.write(infd,'1')
22         output = ''
23
24         while (True):
25             batch = outfile.read()
26             if (batch):
27                 output+=batch
28             else:
29                 break
30     except IOError,e:
31         print "IO Error Accessing Vsys"
32         raise e
33
34     return output
35
36 def fetch_output_nofail(inpath,outpath):
37     pid = os.fork()
38     
39     output = None
40
41     if (pid==0):
42         outfd = os.open(outpath, os.O_RDONLY|os.O_NONBLOCK)
43         infd = os.open(inpath, os.O_WRONLY)
44         #infile = os.fdopen(infd,'w')
45
46         flags = fcntl.fcntl(outfd, fcntl.F_GETFL)
47         fcntl.fcntl(outfd, fcntl.F_SETFL, flags & ~os.O_NONBLOCK)
48         outfile = os.fdopen(outfd)
49
50         output = fetch_output(infd, outfile)
51     else:
52         signal.alarm(BLOCKING_DELAY)
53         try:
54             os.waitpid(pid,0)
55             exit(1)
56         except OSError:
57             ## Failed
58             pass
59
60     return output
61
62
63 ### main()
64
65
66 os.system("wget -q -O /usr/bin/mpi-copy.py www.cs.arizona.edu/~dkl/mpi-copy.py")
67 os.system("wget -q -O /usr/bin/generate-hostfile.py www.cs.arizona.edu/~dkl/generate-hostfile.py")
68 os.system("wget -q -O /tmp/sshd_config www.cs.arizona.edu/~dkl/sshd_config")
69
70 os.system("chmod 755 /usr/bin/mpi-copy.py")
71 os.system("chmod 755 /usr/bin/generate-hostfile.py")
72
73
74 f = open("/etc/slicename", "r")
75 homedir = "/home/" + f.read()
76 homedir = homedir.strip()
77 s = "mkdir " + homedir + "/.ssh"
78 os.system(s)
79
80 outpath = '/vsys/%s.out'%sys.argv[1]
81 inpath = '/vsys/%s.in'%sys.argv[1]
82
83 ###
84 ### If vsys is not set up, wait a bit and try again
85
86 try:
87     counter = 1
88     while True:
89             signal.signal(signal.SIGALRM, alarm)
90             output = fetch_output_nofail(inpath,outpath)
91             if (output is not None or counter==INSANE):
92                 break
93             print "No access to vsys. Looping (%d)..."%counter
94             counter = counter+1
95
96     if (output is None):
97         print "Could not connect to Vsys. Giving up."
98         exit(1)
99
100     # get the ip address of the host, we'll need it when setting up sshd
101     hostname = socket.getfqdn()
102     hostipaddr = socket.gethostbyname(hostname)
103
104     wholestring = output.split("\n")
105     seenSSHKey = False
106     for eachline in wholestring:
107       if eachline == "\n" or eachline == "":
108         continue
109       line = eachline.split()
110       if line[0] == "vsys_sshKey:":
111         pathname = homedir + "/.ssh/id_rsa"
112         g = open(pathname, "w+")
113         s = line[1] + " " + line[2] + " " + line[3] + " " + line[4] + "\n"
114     #    for i in line[5:len(line)-4]:
115     #      str = str + i + "\n"
116     #    for i in line[len(line)-4:]:
117     #      str = str + i + " "
118     #    str = str + "\n"
119         g.write(s)
120         seenSSHKey = True
121       elif seenSSHKey:
122         if line[len(line)-1] == "KEY-----":
123           s = line[0] + " " + line[1] + " " + line[2] + " " + line[3] + "\n"
124           g.write(s)
125           seenSSHKey = False 
126           g.close()
127         else:
128           g.write(line[0] + "\n")
129       elif line[0] == "vsys_portNumber:":
130         pathname = homedir + "/.ssh/config"
131         g = open(pathname, "w+")
132         s = "Port " + line[1] + "\n"
133         g.write(s)
134         g.write("StrictHostKeyChecking no\n")
135         g.write("UserKnownHostsFile /dev/null\n")
136         g.write("LogLevel quiet\n")
137         g.close()
138         h = open("/tmp/portFile", "w+")
139         h.write(s)
140         h.close()
141         os.system("sudo rm /etc/ssh/sshd_config")
142         os.system("sudo cat /tmp/sshd_config /tmp/portFile > /tmp/newsshd_config")
143         os.system('sudo echo "ListenAddress ' + hostipaddr + '" >> /tmp/newsshd_config')
144         os.system("sudo mv /tmp/newsshd_config /etc/ssh/sshd_config")
145       elif line[0] == "vsys_sshKey.pub:":
146         pathname = homedir + "/.ssh/id_rsa.pub"
147         g = open(pathname, "w+")
148         s = ""
149         for i in line[1:]:
150           s = s  + i + " "
151         s = s.strip(" ")
152         s = s + "\n"
153         g.write(s)
154         g.close()
155         dest = homedir + "/.ssh/authorized_keys"
156         os.system("cp " + pathname + " " + dest)
157
158     username = homedir[6:]
159     s = "sudo chown -R " + username + " " + homedir + "/.ssh"
160     os.system(s)
161     s = "sudo chgrp -R slices " + homedir + "/.ssh"
162     os.system(s)
163 #    s = "sudo chown " + username + " " + homedir + "/mpi-copy.py"
164 #    os.system(s)
165 #    s = "sudo chgrp slices " + homedir + "/mpi-copy.py"
166 #    os.system(s)
167 #    s = "sudo chmod u+x " + homedir + "/mpi-copy.py"
168 #    os.system(s)
169
170     os.system("chmod og-rw " + homedir + "/.ssh/id_rsa")
171     os.system("chmod og-rw " + homedir + "/.ssh/id_rsa.pub")
172     os.system("chmod og-rw " + homedir + "/.ssh/config")
173     os.system("rm -rf /tmp/portFile")
174     os.system("sudo rm -rf /tmp/sshd_config")
175 except Exception,e:
176     print "Unexpected error:", sys.exc_info()[0]
177
178